From 5f6a4dbc40bb812eda8790c95a80831e7d9f7ba4 Mon Sep 17 00:00:00 2001 From: Leslie Lam Date: Mon, 10 Jun 2024 14:27:21 -0400 Subject: [PATCH 1/4] [ENTERPRISE-1418] Add support for plain JWT authentication --- dbt/adapters/snowflake/connections.py | 35 ++++++++++++++++++++++----- tests/unit/test_snowflake_adapter.py | 32 ++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py index aca115b4b..99ecf2948 100644 --- a/dbt/adapters/snowflake/connections.py +++ b/dbt/adapters/snowflake/connections.py @@ -43,7 +43,7 @@ from dbt.adapters.sql import SQLConnectionManager from dbt.adapters.events.logging import AdapterLogger from dbt_common.events.functions import warn_or_error -from dbt.adapters.events.types import AdapterEventWarning +from dbt.adapters.events.types import AdapterEventWarning, AdapterEventError from dbt_common.ui import line_wrap_message, warning_tag @@ -70,7 +70,7 @@ class SnowflakeAdapterResponse(AdapterResponse): @dataclass class SnowflakeCredentials(Credentials): account: str - user: str + user: Optional[str] = None warehouse: Optional[str] = None role: Optional[str] = None password: Optional[str] = None @@ -96,15 +96,29 @@ 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.authenticator not in ["oauth", "jwt"]: + if self.token: + warn_or_error( + AdapterEventWarning( + base_msg=( + "The token parameter was set, but the authenticator was " + "not set to 'oauth' or 'jwt'." + ) + ) + ) + + if not self.user: + # The user attribute is only optional if 'authenticator' is 'jwt' or 'oauth' + warn_or_error(AdapterEventError(base_msg="'user' is a required property.")) + self.account = self.account.replace("_", "-") @property @@ -146,6 +160,8 @@ def auth_args(self): # Pull all of the optional authentication args for the connector, # let connector handle the actual arg validation result = {} + if self.user: + result["user"] = self.user if self.password: result["password"] = self.password if self.host: @@ -180,6 +196,14 @@ def auth_args(self): ) result["token"] = token + + elif self.authenticator == "jwt": + # If authenticator is 'jwt', then the 'token' value should be used + # unmodified. We expose this as 'jwt' in the profile, but the value + # passed into the snowflake.connect method should still be 'oauth' + result["token"] = self.token + result["authenticator"] = "oauth" + # enable id token cache for linux result["client_store_temporary_credential"] = True # enable mfa token cache for linux @@ -346,7 +370,6 @@ def connect(): handle = snowflake.connector.connect( account=creds.account, - user=creds.user, database=creds.database, schema=creds.schema, warehouse=creds.warehouse, diff --git a/tests/unit/test_snowflake_adapter.py b/tests/unit/test_snowflake_adapter.py index ff92b9b65..f6a768da8 100644 --- a/tests/unit/test_snowflake_adapter.py +++ b/tests/unit/test_snowflake_adapter.py @@ -550,6 +550,38 @@ 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", user=None + ) + 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", + warehouse="test_warehouse", + authenticator="oauth", + 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" From 1a5216fb82b3087ef0f42fa703ca226e5a8b1835 Mon Sep 17 00:00:00 2001 From: Leslie Lam Date: Mon, 10 Jun 2024 17:10:47 -0400 Subject: [PATCH 2/4] Run changie new --- .changes/unreleased/Features-20240610-171026.yaml | 6 ++++++ dbt/adapters/snowflake/connections.py | 4 +++- 2 files changed, 9 insertions(+), 1 deletion(-) create mode 100644 .changes/unreleased/Features-20240610-171026.yaml diff --git a/.changes/unreleased/Features-20240610-171026.yaml b/.changes/unreleased/Features-20240610-171026.yaml new file mode 100644 index 000000000..5cc055160 --- /dev/null +++ b/.changes/unreleased/Features-20240610-171026.yaml @@ -0,0 +1,6 @@ +kind: Features +body: Support JWT Authentication +time: 2024-06-10T17:10:26.421463-04:00 +custom: + Author: llam15 + Issue: 1079 726 diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py index 99ecf2948..ba79e03d1 100644 --- a/dbt/adapters/snowflake/connections.py +++ b/dbt/adapters/snowflake/connections.py @@ -117,7 +117,9 @@ def __post_init__(self): if not self.user: # The user attribute is only optional if 'authenticator' is 'jwt' or 'oauth' - warn_or_error(AdapterEventError(base_msg="'user' is a required property.")) + warn_or_error( + AdapterEventError(base_msg="Invalid profile: 'user' is a required property.") + ) self.account = self.account.replace("_", "-") From aa21c140d3a8c88dad27d8c8b8d781b2efde7f87 Mon Sep 17 00:00:00 2001 From: Leslie Lam Date: Tue, 11 Jun 2024 13:59:00 -0400 Subject: [PATCH 3/4] wip: functional test for JWT --- tests/functional/oauth/test_jwt.py | 115 +++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 tests/functional/oauth/test_jwt.py diff --git a/tests/functional/oauth/test_jwt.py b/tests/functional/oauth/test_jwt.py new file mode 100644 index 000000000..bd5177b13 --- /dev/null +++ b/tests/functional/oauth/test_jwt.py @@ -0,0 +1,115 @@ +""" +The first time using an account for testing, you should run this: + +``` +CREATE OR REPLACE SECURITY INTEGRATION DBT_INTEGRATION_TEST_OAUTH + TYPE = OAUTH + ENABLED = TRUE + OAUTH_CLIENT = CUSTOM + OAUTH_CLIENT_TYPE = 'CONFIDENTIAL' + OAUTH_REDIRECT_URI = 'http://localhost:8080' + oauth_issue_refresh_tokens = true + OAUTH_ALLOW_NON_TLS_REDIRECT_URI = true + BLOCKED_ROLES_LIST = + oauth_refresh_token_validity = 7776000; +``` + + +Every month (or any amount <90 days): + +Run `select SYSTEM$SHOW_OAUTH_CLIENT_SECRETS('DBT_INTEGRATION_TEST_OAUTH');` + +The only row/column of output should be a json blob, it goes (within single +quotes!) as the second argument to the server script: + +python scripts/werkzeug-refresh-token.py ${acount_name} '${json_blob}' + +Open http://localhost:8080 + +Log in as the test user, get a response page with some environment variables. +Update CI providers and test.env with the new values (If you kept the security +integration the same, just the refresh token changed) +""" + +import pytest +import os +from dbt.tests.util import run_dbt, check_relations_equal + +from dbt.adapters.snowflake import SnowflakeCredentials + +_MODELS__MODEL_1_SQL = """ +select 1 as id +""" + + +_MODELS__MODEL_2_SQL = """ +select 2 as id +""" + + +_MODELS__MODEL_3_SQL = """ +select * from {{ ref('model_1') }} +union all +select * from {{ ref('model_2') }} +""" + + +_MODELS__MODEL_4_SQL = """ +select 1 as id +union all +select 2 as id +""" + + +class TestSnowflakeJWT: + """Tests that setting authenticator: jwt allows setting token to a plain JWT + that will be passed into the Snowflake connection without modification.""" + + @pytest.fixture(scope="class", autouse=True) + def access_token(self): + """Because JWTs are short-lived, we need to get a fresh JWT via the refresh + token flow before running the test. + + This fixture leverages the existing SnowflakeCredentials._get_access_token + method to retrieve a valid JWT from Snowflake. + """ + client_id = os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_ID") + client_secret = os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_SECRET") + refresh_token = os.getenv("SNOWFLAKE_TEST_OAUTH_REFRESH_TOKEN") + + credentials = SnowflakeCredentials( + account=os.getenv("SNOWFLAKE_TEST_ACCOUNT"), + database="", + schema="", + authenticator="oauth", + oauth_client_id=client_id, + oauth_client_secret=client_secret, + token=refresh_token, + ) + + yield credentials._get_access_token() + + @pytest.fixture(scope="class", autouse=True) + def dbt_profile_target(self, access_token): + return { + "type": "snowflake", + "threads": 4, + "account": os.getenv("SNOWFLAKE_TEST_ACCOUNT"), + "token": access_token, + "database": os.getenv("SNOWFLAKE_TEST_DATABASE"), + "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"), + "authenticator": "jwt", + } + + @pytest.fixture(scope="class") + def models(self): + return { + "model_1.sql": _MODELS__MODEL_1_SQL, + "model_2.sql": _MODELS__MODEL_2_SQL, + "model_3.sql": _MODELS__MODEL_3_SQL, + "model_4.sql": _MODELS__MODEL_4_SQL, + } + + def test_snowflake_basic(self, project): + run_dbt() + check_relations_equal(project.adapter, ["MODEL_3", "MODEL_4"]) From 6bc54be7b0f986d391f2d04d38d7f2a3e3a2fd20 Mon Sep 17 00:00:00 2001 From: Leslie Lam Date: Tue, 11 Jun 2024 14:17:09 -0400 Subject: [PATCH 4/4] clean up, and add some comments --- tests/functional/oauth/test_jwt.py | 38 ++++++------------------------ 1 file changed, 7 insertions(+), 31 deletions(-) diff --git a/tests/functional/oauth/test_jwt.py b/tests/functional/oauth/test_jwt.py index bd5177b13..fbe8e20e6 100644 --- a/tests/functional/oauth/test_jwt.py +++ b/tests/functional/oauth/test_jwt.py @@ -1,34 +1,6 @@ """ -The first time using an account for testing, you should run this: - -``` -CREATE OR REPLACE SECURITY INTEGRATION DBT_INTEGRATION_TEST_OAUTH - TYPE = OAUTH - ENABLED = TRUE - OAUTH_CLIENT = CUSTOM - OAUTH_CLIENT_TYPE = 'CONFIDENTIAL' - OAUTH_REDIRECT_URI = 'http://localhost:8080' - oauth_issue_refresh_tokens = true - OAUTH_ALLOW_NON_TLS_REDIRECT_URI = true - BLOCKED_ROLES_LIST = - oauth_refresh_token_validity = 7776000; -``` - - -Every month (or any amount <90 days): - -Run `select SYSTEM$SHOW_OAUTH_CLIENT_SECRETS('DBT_INTEGRATION_TEST_OAUTH');` - -The only row/column of output should be a json blob, it goes (within single -quotes!) as the second argument to the server script: - -python scripts/werkzeug-refresh-token.py ${acount_name} '${json_blob}' - -Open http://localhost:8080 - -Log in as the test user, get a response page with some environment variables. -Update CI providers and test.env with the new values (If you kept the security -integration the same, just the refresh token changed) +Please follow the instructions in test_oauth.py for instructions on how to set up +the security integration required to retrieve a JWT from Snowflake. """ import pytest @@ -91,14 +63,18 @@ def access_token(self): @pytest.fixture(scope="class", autouse=True) def dbt_profile_target(self, access_token): + """A dbt_profile that has authenticator set to JWT, and token set to + a JWT accepted by Snowflake. Also omits the user, as the user attribute + is optional when the authenticator is set to JWT. + """ return { "type": "snowflake", "threads": 4, "account": os.getenv("SNOWFLAKE_TEST_ACCOUNT"), - "token": access_token, "database": os.getenv("SNOWFLAKE_TEST_DATABASE"), "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"), "authenticator": "jwt", + "token": access_token, } @pytest.fixture(scope="class")