From 2576a08379e248d393b9a698e58fac0d32c23648 Mon Sep 17 00:00:00 2001
From: Mike Alfare <13974384+mikealfare@users.noreply.github.com>
Date: Mon, 24 Jun 2024 13:14:35 -0400
Subject: [PATCH] [Bug] Silent SSO token expiration (#1072)

* changelog

* add an integration test for token expiration

* catch missing access_token key error

* update test to connect initially, then swap out token for a invalid one

* modify test and how we setup profile

* attempt at unit test

* remove functional test attempt

* add name to changelog

---------

Co-authored-by: McKnight-42 <matthew.mcknight@dbtlabs.com>
Co-authored-by: Colin Rogers <111200756+colin-rogers-dbt@users.noreply.github.com>
Co-authored-by: Matthew McKnight <91097623+McKnight-42@users.noreply.github.com>
---
 .../unreleased/Fixes-20240605-125611.yaml     |  6 ++++
 dbt/adapters/snowflake/connections.py         |  6 +++-
 tests/functional/oauth/test_oauth.py          |  4 +--
 tests/unit/test_connections.py                | 33 ++++++++++++++++++-
 4 files changed, 45 insertions(+), 4 deletions(-)
 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..c4560774c
--- /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, McKnight-42
+  Issue: "851"
diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py
index 4db007f19..6e9a5aaba 100644
--- a/dbt/adapters/snowflake/connections.py
+++ b/dbt/adapters/snowflake/connections.py
@@ -266,7 +266,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):
diff --git a/tests/functional/oauth/test_oauth.py b/tests/functional/oauth/test_oauth.py
index 89daece0f..c8986763e 100644
--- a/tests/functional/oauth/test_oauth.py
+++ b/tests/functional/oauth/test_oauth.py
@@ -31,9 +31,9 @@
 integration the same, just the refresh token changed)
 """
 
-import pytest
 import os
-from dbt.tests.util import run_dbt, check_relations_equal
+from dbt.tests.util import check_relations_equal, run_dbt
+import pytest
 
 
 _MODELS__MODEL_1_SQL = """
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()