From be5717ea325162367864b1f940bdf5dc9cf1554d Mon Sep 17 00:00:00 2001
From: Pavel Grivachev <pavel@fivetran.com>
Date: Wed, 1 Nov 2023 16:33:04 +0100
Subject: [PATCH] support snowflake mfa (#1326)

* support mfa for snowflake client

* bump version
---
 package-lock.json                                 |  4 ++--
 package.json                                      |  2 +-
 server/src/snowflake/SnowflakeKeyPairProfile.ts   |  4 +++-
 server/src/snowflake/SnowflakeUserPassProfile.ts  | 15 +++++++++++++--
 server/src/test/profiles/snowflake.yml            | 15 +++++++++++++++
 .../snowflake/SnowflakeUserPassProfile.spec.ts    |  4 ++++
 6 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/package-lock.json b/package-lock.json
index 797aa77d..248e0841 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,12 +1,12 @@
 {
   "name": "dbt-language-server",
-  "version": "0.33.0",
+  "version": "0.33.1",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "dbt-language-server",
-      "version": "0.33.0",
+      "version": "0.33.1",
       "hasInstallScript": true,
       "license": "MIT",
       "devDependencies": {
diff --git a/package.json b/package.json
index 328f2c1f..58bede65 100644
--- a/package.json
+++ b/package.json
@@ -3,7 +3,7 @@
   "displayName": "Wizard for dbt Core (TM)",
   "description": "This extension will help you work with dbt",
   "icon": "images/Icon.png",
-  "version": "0.33.0",
+  "version": "0.33.1",
   "publisher": "Fivetran",
   "license": "MIT",
   "preview": true,
diff --git a/server/src/snowflake/SnowflakeKeyPairProfile.ts b/server/src/snowflake/SnowflakeKeyPairProfile.ts
index 9757622d..53ce135c 100644
--- a/server/src/snowflake/SnowflakeKeyPairProfile.ts
+++ b/server/src/snowflake/SnowflakeKeyPairProfile.ts
@@ -17,6 +17,7 @@ export class SnowflakeKeyPairProfile implements DbtProfile {
     database: z.string(),
     warehouse: z.string(),
     schema: z.string(),
+    client_session_keep_alive: z.optional(z.boolean()),
   });
 
   getDocsUrl(): string {
@@ -46,8 +47,9 @@ export class SnowflakeKeyPairProfile implements DbtProfile {
       warehouse,
       database,
       schema,
+      client_session_keep_alive: clientSessionKeepAlive,
     } = profile;
-    const connection = createConnection({ account, username, privateKeyPath, privateKeyPass, warehouse, database, schema });
+    const connection = createConnection({ account, username, privateKeyPath, privateKeyPass, warehouse, database, schema, clientSessionKeepAlive });
 
     const client = new SnowflakeClient(database, connection);
 
diff --git a/server/src/snowflake/SnowflakeUserPassProfile.ts b/server/src/snowflake/SnowflakeUserPassProfile.ts
index 73d1bf8e..1f276c96 100644
--- a/server/src/snowflake/SnowflakeUserPassProfile.ts
+++ b/server/src/snowflake/SnowflakeUserPassProfile.ts
@@ -17,6 +17,8 @@ export class SnowflakeUserPassProfile implements DbtProfile {
     database: z.string(),
     warehouse: z.string(),
     schema: z.string(),
+    authenticator: z.optional(z.string()),
+    client_session_keep_alive: z.optional(z.boolean()),
   });
 
   getDocsUrl(): string {
@@ -38,8 +40,17 @@ export class SnowflakeUserPassProfile implements DbtProfile {
   }
 
   private async createClientInternal(profile: z.infer<typeof SnowflakeUserPassProfile.SCHEMA>): Promise<Result<DbtDestinationClient, string>> {
-    const { account, user: username, password, warehouse, database, schema } = profile;
-    const connection = createConnection({ account, username, password, warehouse, database, schema });
+    const {
+      account,
+      user: username,
+      password,
+      warehouse,
+      database,
+      schema,
+      authenticator,
+      client_session_keep_alive: clientSessionKeepAlive,
+    } = profile;
+    const connection = createConnection({ account, username, password, warehouse, database, schema, authenticator, clientSessionKeepAlive });
 
     const client = new SnowflakeClient(database, connection);
 
diff --git a/server/src/test/profiles/snowflake.yml b/server/src/test/profiles/snowflake.yml
index b7e1fae7..18bbb807 100644
--- a/server/src/test/profiles/snowflake.yml
+++ b/server/src/test/profiles/snowflake.yml
@@ -11,6 +11,21 @@ correct_user_password:
       warehouse: warehouse_name
       schema: dbt_schema
 
+correct_user_password_mfa:
+  target: dev
+  outputs:
+    dev:
+      type: snowflake
+      account: acc
+      user: username
+      password: password
+      role: role
+      database: database_name
+      warehouse: warehouse_name
+      schema: dbt_schema
+      authenticator: externalbrowser
+      client_session_keep_alive: True
+
 user_password_missing_user:
   target: dev
   outputs:
diff --git a/server/src/test/snowflake/SnowflakeUserPassProfile.spec.ts b/server/src/test/snowflake/SnowflakeUserPassProfile.spec.ts
index 37bddf3a..e6cbaea2 100644
--- a/server/src/test/snowflake/SnowflakeUserPassProfile.spec.ts
+++ b/server/src/test/snowflake/SnowflakeUserPassProfile.spec.ts
@@ -7,6 +7,10 @@ describe('SnowflakeUserPassProfile', () => {
     shouldPassValidProfile(SNOWFLAKE_CONFIG, 'correct_user_password');
   });
 
+  it('Should pass valid profile with MFA', () => {
+    shouldPassValidProfile(SNOWFLAKE_CONFIG, 'correct_user_password_mfa');
+  });
+
   it('Should require user', () => {
     shouldRequireField('user');
   });