From 1835c02399a651833170438cd75257e6afe69aac Mon Sep 17 00:00:00 2001 From: Mike Alfare Date: Wed, 5 Jun 2024 12:56:34 -0400 Subject: [PATCH 1/8] changelog --- .changes/unreleased/Fixes-20240605-125611.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 .changes/unreleased/Fixes-20240605-125611.yaml diff --git a/.changes/unreleased/Fixes-20240605-125611.yaml b/.changes/unreleased/Fixes-20240605-125611.yaml new file mode 100644 index 000000000..c96f5742f --- /dev/null +++ b/.changes/unreleased/Fixes-20240605-125611.yaml @@ -0,0 +1,6 @@ +kind: Fixes +body: Surface SSO token expiration in logs +time: 2024-06-05T12:56:11.802237-04:00 +custom: + Author: mikealfare + Issue: "851" From 3704a155de813706d06a9ba22d813d0c00d87cac Mon Sep 17 00:00:00 2001 From: Mike Alfare Date: Wed, 5 Jun 2024 13:28:41 -0400 Subject: [PATCH 2/8] add an integration test for token expiration --- tests/functional/oauth/test_oauth.py | 31 ++++++++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/tests/functional/oauth/test_oauth.py b/tests/functional/oauth/test_oauth.py index 89daece0f..ff2bef714 100644 --- a/tests/functional/oauth/test_oauth.py +++ b/tests/functional/oauth/test_oauth.py @@ -31,9 +31,11 @@ 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.exceptions.connection import FailedToConnectError +from dbt.tests.util import check_relations_equal, run_dbt +import pytest _MODELS__MODEL_1_SQL = """ @@ -88,3 +90,28 @@ def models(self): def test_snowflake_basic(self, project): run_dbt() check_relations_equal(project.adapter, ["MODEL_3", "MODEL_4"]) + + +class TestSnowflakeOAuthExpiration: + @pytest.fixture(scope="class") + def dbt_profile_target(self): + return { + "type": "snowflake", + "threads": 4, + "account": os.getenv("SNOWFLAKE_TEST_ACCOUNT"), + "user": os.getenv("SNOWFLAKE_TEST_USER"), + "oauth_client_id": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_ID"), + "oauth_client_secret": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_SECRET"), + "token": "THIS_TOKEN_DOES_NOT_EXIST", + "database": os.getenv("SNOWFLAKE_TEST_DATABASE"), + "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"), + "authenticator": "oauth", + } + + @pytest.fixture(scope="class") + def models(self): + return {"model_1.sql": _MODELS__MODEL_1_SQL} + + def test_token_expiration(self, project): + with pytest.raises(FailedToConnectError): + run_dbt(["run"], expect_pass=False) From 667bcb6da70e2e4fcd88ecb3779eb4f69f8ae8e2 Mon Sep 17 00:00:00 2001 From: Mike Alfare Date: Wed, 5 Jun 2024 13:48:00 -0400 Subject: [PATCH 3/8] catch missing access_token key error --- dbt/adapters/snowflake/connections.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py index 1d6e31c93..ab330a2ac 100644 --- a/dbt/adapters/snowflake/connections.py +++ b/dbt/adapters/snowflake/connections.py @@ -238,7 +238,11 @@ def _get_access_token(self) -> str: f"""Did not receive valid json with access_token. Showing json response: {result_json}""" ) - + elif "access_token" not in result_json: + raise FailedToConnectError( + "This error occurs when authentication has expired. " + "Please reauth with your auth provider." + ) return result_json["access_token"] def _get_private_key(self): From 502d4489a7cfa1a5d10efed8b810c5809b945071 Mon Sep 17 00:00:00 2001 From: McKnight-42 Date: Fri, 14 Jun 2024 13:53:43 -0500 Subject: [PATCH 4/8] update test to connect initially, then swap out token for a invalid one --- tests/functional/oauth/test_oauth.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/functional/oauth/test_oauth.py b/tests/functional/oauth/test_oauth.py index ff2bef714..618486809 100644 --- a/tests/functional/oauth/test_oauth.py +++ b/tests/functional/oauth/test_oauth.py @@ -32,9 +32,9 @@ """ import os - +import yaml from dbt.adapters.exceptions.connection import FailedToConnectError -from dbt.tests.util import check_relations_equal, run_dbt +from dbt.tests.util import check_relations_equal, run_dbt, write_file import pytest @@ -93,7 +93,7 @@ def test_snowflake_basic(self, project): class TestSnowflakeOAuthExpiration: - @pytest.fixture(scope="class") + @pytest.fixture(scope="class", autouse=True) def dbt_profile_target(self): return { "type": "snowflake", @@ -102,7 +102,7 @@ def dbt_profile_target(self): "user": os.getenv("SNOWFLAKE_TEST_USER"), "oauth_client_id": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_ID"), "oauth_client_secret": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_SECRET"), - "token": "THIS_TOKEN_DOES_NOT_EXIST", + "token": os.getenv("SNOWFLAKE_TEST_OAUTH_REFRESH_TOKEN"), "database": os.getenv("SNOWFLAKE_TEST_DATABASE"), "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"), "authenticator": "oauth", @@ -112,6 +112,10 @@ def dbt_profile_target(self): def models(self): return {"model_1.sql": _MODELS__MODEL_1_SQL} - def test_token_expiration(self, project): + def test_token_expiration(self, project, dbt_profile_data): + # change token which is included in the connection_info + dbt_profile_data["test"]["outputs"]["default"]["token"] = "THIS_TOKEN_DOES_NOT_EXIST" + write_file(yaml.safe_dump(dbt_profile_data), project.profiles_dir, "profiles.yml") + with pytest.raises(FailedToConnectError): run_dbt(["run"], expect_pass=False) From 5664c51670c0dd1fc28188ec227d8db90b052ffc Mon Sep 17 00:00:00 2001 From: McKnight-42 Date: Fri, 14 Jun 2024 16:14:08 -0500 Subject: [PATCH 5/8] modify test and how we setup profile --- tests/functional/oauth/test_oauth.py | 35 +++++++++++++++------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/tests/functional/oauth/test_oauth.py b/tests/functional/oauth/test_oauth.py index 618486809..97847770b 100644 --- a/tests/functional/oauth/test_oauth.py +++ b/tests/functional/oauth/test_oauth.py @@ -93,29 +93,32 @@ def test_snowflake_basic(self, project): class TestSnowflakeOAuthExpiration: + profile = { + "type": "snowflake", + "threads": 4, + "account": os.getenv("SNOWFLAKE_TEST_ACCOUNT"), + "user": os.getenv("SNOWFLAKE_TEST_USER"), + "oauth_client_id": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_ID"), + "oauth_client_secret": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_SECRET"), + "token": os.getenv("SNOWFLAKE_TEST_OAUTH_REFRESH_TOKEN"), + "database": os.getenv("SNOWFLAKE_TEST_DATABASE"), + "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"), + "authenticator": "oauth", + } + @pytest.fixture(scope="class", autouse=True) def dbt_profile_target(self): - return { - "type": "snowflake", - "threads": 4, - "account": os.getenv("SNOWFLAKE_TEST_ACCOUNT"), - "user": os.getenv("SNOWFLAKE_TEST_USER"), - "oauth_client_id": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_ID"), - "oauth_client_secret": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_SECRET"), - "token": os.getenv("SNOWFLAKE_TEST_OAUTH_REFRESH_TOKEN"), - "database": os.getenv("SNOWFLAKE_TEST_DATABASE"), - "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"), - "authenticator": "oauth", - } + return self.profile @pytest.fixture(scope="class") def models(self): return {"model_1.sql": _MODELS__MODEL_1_SQL} - def test_token_expiration(self, project, dbt_profile_data): + def test_token_expiration(self, project): # change token which is included in the connection_info - dbt_profile_data["test"]["outputs"]["default"]["token"] = "THIS_TOKEN_DOES_NOT_EXIST" - write_file(yaml.safe_dump(dbt_profile_data), project.profiles_dir, "profiles.yml") + project.adapter.config.credentials.token = "THIS_TOKEN_DOES_NOT_EXIST" + + self.profile["token"] = "THIS_TOKEN_DOES_NOT_EXIST" with pytest.raises(FailedToConnectError): - run_dbt(["run"], expect_pass=False) + run_dbt(["run", "--profile", self.profile], expect_pass=False) From 690626b919a51c715e778a1b3ba109e6c8d00a2b Mon Sep 17 00:00:00 2001 From: McKnight-42 Date: Mon, 24 Jun 2024 10:38:42 -0500 Subject: [PATCH 6/8] attempt at unit test --- tests/unit/test_connections.py | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_connections.py b/tests/unit/test_connections.py index dd452b3cb..fb9c57615 100644 --- a/tests/unit/test_connections.py +++ b/tests/unit/test_connections.py @@ -1,6 +1,9 @@ import os +import pytest from importlib import reload -from unittest.mock import Mock +from unittest.mock import Mock, patch +import multiprocessing +from dbt.adapters.exceptions.connection import FailedToConnectError import dbt.adapters.snowflake.connections as connections import dbt.adapters.events.logging @@ -36,3 +39,31 @@ def test_connnections_credentials_replaces_underscores_with_hyphens(): } creds = connections.SnowflakeCredentials(**credentials) assert creds.account == "account-id-with-underscores" + + +def test_snowflake_oauth_expired_token_raises_error(): + credentials = { + "account": "test_account", + "user": "test_user", + "authenticator": "oauth", + "token": "expired_or_incorrect_token", + "database": "database", + "schema": "schema", + } + + mp_context = multiprocessing.get_context("spawn") + mock_credentials = connections.SnowflakeCredentials(**credentials) + + with patch.object( + connections.SnowflakeConnectionManager, + "open", + side_effect=FailedToConnectError( + "This error occurs when authentication has expired. " + "Please reauth with your auth provider." + ), + ): + + adapter = connections.SnowflakeConnectionManager(mock_credentials, mp_context) + + with pytest.raises(FailedToConnectError): + adapter.open() From 5246386e6ee9178761af6ffa6095627f382286da Mon Sep 17 00:00:00 2001 From: McKnight-42 Date: Mon, 24 Jun 2024 10:40:37 -0500 Subject: [PATCH 7/8] remove functional test attempt --- tests/functional/oauth/test_oauth.py | 36 +--------------------------- 1 file changed, 1 insertion(+), 35 deletions(-) diff --git a/tests/functional/oauth/test_oauth.py b/tests/functional/oauth/test_oauth.py index 97847770b..c8986763e 100644 --- a/tests/functional/oauth/test_oauth.py +++ b/tests/functional/oauth/test_oauth.py @@ -32,9 +32,7 @@ """ import os -import yaml -from dbt.adapters.exceptions.connection import FailedToConnectError -from dbt.tests.util import check_relations_equal, run_dbt, write_file +from dbt.tests.util import check_relations_equal, run_dbt import pytest @@ -90,35 +88,3 @@ def models(self): def test_snowflake_basic(self, project): run_dbt() check_relations_equal(project.adapter, ["MODEL_3", "MODEL_4"]) - - -class TestSnowflakeOAuthExpiration: - profile = { - "type": "snowflake", - "threads": 4, - "account": os.getenv("SNOWFLAKE_TEST_ACCOUNT"), - "user": os.getenv("SNOWFLAKE_TEST_USER"), - "oauth_client_id": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_ID"), - "oauth_client_secret": os.getenv("SNOWFLAKE_TEST_OAUTH_CLIENT_SECRET"), - "token": os.getenv("SNOWFLAKE_TEST_OAUTH_REFRESH_TOKEN"), - "database": os.getenv("SNOWFLAKE_TEST_DATABASE"), - "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"), - "authenticator": "oauth", - } - - @pytest.fixture(scope="class", autouse=True) - def dbt_profile_target(self): - return self.profile - - @pytest.fixture(scope="class") - def models(self): - return {"model_1.sql": _MODELS__MODEL_1_SQL} - - def test_token_expiration(self, project): - # change token which is included in the connection_info - project.adapter.config.credentials.token = "THIS_TOKEN_DOES_NOT_EXIST" - - self.profile["token"] = "THIS_TOKEN_DOES_NOT_EXIST" - - with pytest.raises(FailedToConnectError): - run_dbt(["run", "--profile", self.profile], expect_pass=False) From 13e92677b829c841ec41496509f48c595d6559ee Mon Sep 17 00:00:00 2001 From: McKnight-42 Date: Mon, 24 Jun 2024 12:00:08 -0500 Subject: [PATCH 8/8] add name to changelog --- .changes/unreleased/Fixes-20240605-125611.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.changes/unreleased/Fixes-20240605-125611.yaml b/.changes/unreleased/Fixes-20240605-125611.yaml index c96f5742f..c4560774c 100644 --- a/.changes/unreleased/Fixes-20240605-125611.yaml +++ b/.changes/unreleased/Fixes-20240605-125611.yaml @@ -2,5 +2,5 @@ kind: Fixes body: Surface SSO token expiration in logs time: 2024-06-05T12:56:11.802237-04:00 custom: - Author: mikealfare + Author: mikealfare, McKnight-42 Issue: "851"