From e1fd7b1d93f463ad6f2d3399be23cbb174dc6f2e Mon Sep 17 00:00:00 2001
From: Leslie Lam <leslie.lam@dbtlabs.com>
Date: Mon, 10 Jun 2024 14:27:21 -0400
Subject: [PATCH] [ENTERPRISE-1418] Add support for plain JWT authentication

---
 dbt/adapters/snowflake/connections.py | 21 ++++++++++++++---
 tests/unit/test_snowflake_adapter.py  | 34 +++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py
index aca115b4b..4d702e894 100644
--- a/dbt/adapters/snowflake/connections.py
+++ b/dbt/adapters/snowflake/connections.py
@@ -96,15 +96,24 @@ class SnowflakeCredentials(Credentials):
     reuse_connections: Optional[bool] = None
 
     def __post_init__(self):
-        if self.authenticator != "oauth" and (
-            self.oauth_client_secret or self.oauth_client_id or self.token
-        ):
+        if self.authenticator != "oauth" and (self.oauth_client_secret or self.oauth_client_id):
             # the user probably forgot to set 'authenticator' like I keep doing
             warn_or_error(
                 AdapterEventWarning(
                     base_msg="Authenticator is not set to oauth, but an oauth-only parameter is set! Did you mean to set authenticator: oauth?"
                 )
             )
+
+        if self.token and self.authenticator not in ["oauth", "jwt"]:
+            warn_or_error(
+                AdapterEventWarning(
+                    base_msg=(
+                        "The token parameter was set, but the authenticator was "
+                        "not set to 'oauth' or 'jwt'."
+                    )
+                )
+            )
+
         self.account = self.account.replace("_", "-")
 
     @property
@@ -180,6 +189,12 @@ def auth_args(self):
                     )
 
                 result["token"] = token
+
+            elif self.authenticator == "jwt":
+                # If authenticator is 'jwt', then the 'token' value should be used
+                # unmodified.
+                result["token"] = self.token
+
             # enable id token cache for linux
             result["client_store_temporary_credential"] = True
             # enable mfa token cache for linux
diff --git a/tests/unit/test_snowflake_adapter.py b/tests/unit/test_snowflake_adapter.py
index ff92b9b65..ab726c916 100644
--- a/tests/unit/test_snowflake_adapter.py
+++ b/tests/unit/test_snowflake_adapter.py
@@ -550,6 +550,40 @@ def test_authenticator_private_key_authentication_no_passphrase(self, mock_get_p
             ]
         )
 
+    def test_authenticator_jwt_authentication(self):
+        self.config.credentials = self.config.credentials.replace(
+            authenticator="jwt",
+            token="my-jwt-token",
+        )
+        self.adapter = SnowflakeAdapter(self.config, get_context("spawn"))
+        conn = self.adapter.connections.set_connection_name(name="new_connection_with_new_config")
+
+        self.snowflake.assert_not_called()
+        conn.handle
+        self.snowflake.assert_has_calls(
+            [
+                mock.call(
+                    account="test-account",
+                    autocommit=True,
+                    client_session_keep_alive=False,
+                    database="test_database",
+                    role=None,
+                    schema="public",
+                    user="test_user",
+                    warehouse="test_warehouse",
+                    authenticator="jwt",
+                    token="my-jwt-token",
+                    private_key=None,
+                    application="dbt",
+                    client_request_mfa_token=True,
+                    client_store_temporary_credential=True,
+                    insecure_mode=False,
+                    session_parameters={},
+                    reuse_connections=None,
+                )
+            ]
+        )
+
     def test_query_tag(self):
         self.config.credentials = self.config.credentials.replace(
             password="test_password", query_tag="test_query_tag"