feat(snowflake): Add support for OAuth 2.0 authentication (#36856)

Co-authored-by: Evan Rusackas <evan@preset.io>
Co-authored-by: Evan Rusackas <evan@rusackas.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Joe Li <joe@preset.io>

tests/unit_tests/sql_lab_test.py

@@ -465,86 +465,93 @@ def test_get_sql_results_oauth2(mocker: MockerFixture, app) -> None:
     """
     Test that `get_sql_results` works with OAuth2.
     """
+    # Pushed/popped manually (rather than via a ``with`` block) so the
+    # ``finally`` below still pops it if an assertion fails, preventing the
+    # request context from leaking into later tests in the same session.
     app_context = app.test_request_context()
     app_context.push()
 
-    mocker.patch(
-        "superset.db_engine_specs.base.uuid4",
-        return_value=UUID("fb11f528-6eba-4a8a-837e-6b0d39ee9187"),
-    )
-    mocker.patch(
-        "superset.db_engine_specs.base.generate_code_verifier",
-        return_value="xkBPVZoFChVcy3VZ2l5u7d0FZPTU-olO7HtsAOok2IUGigyoZ62tG_oldy2xg9_HdqPKrWUmKZLmU-CUqz_SQ",
-    )
-    mocker.patch("superset.daos.key_value.KeyValueDAO.delete_expired_entries")
-    mocker.patch("superset.daos.key_value.KeyValueDAO.create_entry")
-    mocker.patch("superset.db_engine_specs.base.db.session.commit")
+    try:
+        mocker.patch(
+            "superset.db_engine_specs.base.uuid4",
+            return_value=UUID("fb11f528-6eba-4a8a-837e-6b0d39ee9187"),
+        )
+        mocker.patch(
+            "superset.db_engine_specs.base.generate_code_verifier",
+            return_value="xkBPVZoFChVcy3VZ2l5u7d0FZPTU-olO7HtsAOok2IUGigyoZ62tG_oldy2xg9_HdqPKrWUmKZLmU-CUqz_SQ",
+        )
+        mocker.patch("superset.daos.key_value.KeyValueDAO.delete_expired_entries")
+        mocker.patch("superset.daos.key_value.KeyValueDAO.create_entry")
+        mocker.patch("superset.db_engine_specs.base.db.session.commit")
 
-    g = mocker.patch("superset.db_engine_specs.base.g")
-    g.user = mocker.MagicMock()
-    g.user.id = 42
+        g = mocker.patch("superset.db_engine_specs.base.g")
+        g.user = mocker.MagicMock()
+        g.user.id = 42
 
-    database = Database(
-        id=1,
-        database_name="my_db",
-        sqlalchemy_uri="sqlite://",
-        encrypted_extra=json.dumps(oauth2_client_info),
-    )
-    database.db_engine_spec.oauth2_exception = OAuth2Error
-    get_sqla_engine = mocker.patch.object(database, "get_sqla_engine")
-    get_sqla_engine().__enter__().raw_connection.side_effect = OAuth2Error(
-        "OAuth2 required"
-    )
+        database = Database(
+            id=1,
+            database_name="my_db",
+            sqlalchemy_uri="sqlite://",
+            encrypted_extra=json.dumps(oauth2_client_info),
+        )
+        database.db_engine_spec.oauth2_exception = OAuth2Error
+        get_sqla_engine = mocker.patch.object(database, "get_sqla_engine")
+        get_sqla_engine().__enter__().raw_connection.side_effect = OAuth2Error(
+            "OAuth2 required"
+        )
 
-    # `limit` and `select_as_cta_used` must match the real `Query` model's
-    # defaults (nullable Integer -> None, Boolean default=False) so that
-    # `apply_limit` -- called unconditionally before the mocked OAuth2 error
-    # is ever reached -- doesn't try to compare an unconfigured MagicMock
-    # against an int.
-    query = mocker.MagicMock(
-        select_as_cta=False,
-        select_as_cta_used=False,
-        limit=None,
-        database=database,
-    )
-    mocker.patch("superset.sql_lab.get_query", return_value=query)
+        # `limit` and `select_as_cta_used` must match the real `Query` model's
+        # defaults (nullable Integer -> None, Boolean default=False) so that
+        # `apply_limit` -- called unconditionally before the mocked OAuth2 error
+        # is ever reached -- doesn't try to compare an unconfigured MagicMock
+        # against an int.
+        query = mocker.MagicMock(
+            select_as_cta=False,
+            select_as_cta_used=False,
+            limit=None,
+            database=database,
+        )
+        mocker.patch("superset.sql_lab.get_query", return_value=query)
 
-    payload = get_sql_results(query_id=1, rendered_query="SELECT 1")
-    assert payload["status"] == QueryStatus.FAILED
-    assert payload["error"] == "You don't have permission to access the data."
-    assert len(payload["errors"]) == 1
+        payload = get_sql_results(query_id=1, rendered_query="SELECT 1")
+        assert payload["status"] == QueryStatus.FAILED
+        assert payload["error"] == "You don't have permission to access the data."
+        assert len(payload["errors"]) == 1
 
-    error = payload["errors"][0]
-    assert error["message"] == "You don't have permission to access the data."
-    assert error["error_type"] == SupersetErrorType.OAUTH2_REDIRECT
-    assert error["level"] == ErrorLevel.WARNING
-    assert error["extra"]["tab_id"] == "fb11f528-6eba-4a8a-837e-6b0d39ee9187"
-    assert (
-        error["extra"]["redirect_uri"] == "http://example.com/api/v1/database/oauth2/"
-    )
+        error = payload["errors"][0]
+        assert error["message"] == "You don't have permission to access the data."
+        assert error["error_type"] == SupersetErrorType.OAUTH2_REDIRECT
+        assert error["level"] == ErrorLevel.WARNING
+        assert error["extra"]["tab_id"] == "fb11f528-6eba-4a8a-837e-6b0d39ee9187"
+        assert (
+            error["extra"]["redirect_uri"]
+            == "http://example.com/api/v1/database/oauth2/"
+        )
 
-    # Parse the OAuth2 authorization URL and verify components individually,
-    # since the JWT state and PKCE code_challenge are computed deterministically
-    # from mocked inputs but their exact encoding depends on library internals.
-    url = urlparse(error["extra"]["url"])
-    assert url.scheme == "https"
-    assert url.netloc == "abcd1234.snowflakecomputing.com"
-    assert url.path == "/oauth/authorize"
+        # Parse the OAuth2 authorization URL and verify components individually,
+        # since the JWT state and PKCE code_challenge are computed deterministically
+        # from mocked inputs but their exact encoding depends on library internals.
+        url = urlparse(error["extra"]["url"])
+        assert url.scheme == "https"
+        assert url.netloc == "abcd1234.snowflakecomputing.com"
+        assert url.path == "/oauth/authorize"
 
-    params = parse_qs(url.query)
-    assert params["scope"] == ["refresh_token session:role:USERADMIN"]
-    assert params["response_type"] == ["code"]
-    assert params["redirect_uri"] == ["http://example.com/api/v1/database/oauth2/"]
-    assert params["client_id"] == ["my_client_id"]
-    assert params["code_challenge_method"] == ["S256"]
+        params = parse_qs(url.query)
+        assert params["scope"] == ["refresh_token session:role:USERADMIN"]
+        assert params["response_type"] == ["code"]
+        assert params["redirect_uri"] == ["http://example.com/api/v1/database/oauth2/"]
+        assert params["client_id"] == ["my_client_id"]
+        assert params["code_challenge_method"] == ["S256"]
 
-    # Verify PKCE code_challenge matches the mocked code_verifier
-    from superset.utils.oauth2 import generate_code_challenge
+        # Verify PKCE code_challenge matches the mocked code_verifier
+        from superset.utils.oauth2 import generate_code_challenge
 
-    expected_code_challenge = generate_code_challenge(
-        "xkBPVZoFChVcy3VZ2l5u7d0FZPTU-olO7HtsAOok2IUGigyoZ62tG_oldy2xg9_HdqPKrWUmKZLmU-CUqz_SQ"
-    )
-    assert params["code_challenge"] == [expected_code_challenge]
+        expected_code_challenge = generate_code_challenge(
+            "xkBPVZoFChVcy3VZ2l5u7d0FZPTU-olO7HtsAOok2IUGigyoZ62tG_oldy2xg9_HdqPKrWUmKZLmU-CUqz_SQ"
+        )
+        assert params["code_challenge"] == [expected_code_challenge]
+    finally:
+        app_context.pop()
 
 
 def test_apply_rls(mocker: MockerFixture) -> None: