Skip to content

Commit

Permalink
[ENTERPRISE-1418] Add support for plain JWT authentication
Browse files Browse the repository at this point in the history
  • Loading branch information
llam15 committed Jun 10, 2024
1 parent 4480734 commit e1fd7b1
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 3 deletions.
21 changes: 18 additions & 3 deletions dbt/adapters/snowflake/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
34 changes: 34 additions & 0 deletions tests/unit/test_snowflake_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down

0 comments on commit e1fd7b1

Please sign in to comment.