diff --git a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx index 60038e5f911..cb5931bc58d 100644 --- a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx +++ b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx @@ -522,7 +522,7 @@ const ExtraOptions = ({ onChange={onInputChange} > {t( - 'Impersonate logged in user (Presto, Trino, Drill, Hive, Databricks, and Google Sheets)', + 'Impersonate logged in user (Presto, Trino, Drill, Hive, Databricks, Snowflake and Google Sheets)', )} diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index 6025bb7adc0..c0641ad3334 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -20,21 +20,23 @@ import logging import re from datetime import datetime from re import Pattern -from typing import Any, Callable, Optional, TYPE_CHECKING, TypedDict +from typing import Any, Callable, cast, Optional, TYPE_CHECKING, TypedDict from urllib import parse from apispec import APISpec from apispec.ext.marshmallow import MarshmallowPlugin from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization -from flask import current_app as app +from flask import current_app as app, has_request_context from flask_babel import gettext as __ from marshmallow import fields, Schema from sqlalchemy import text, types from sqlalchemy.engine.reflection import Inspector from sqlalchemy.engine.url import URL +from sqlalchemy.exc import DatabaseError as SqlalchemyDatabaseError from sqlalchemy.sql.elements import ColumnElement +from superset import is_feature_enabled, security_manager from superset.constants import TimeGrain from superset.databases.utils import make_url_safe from superset.db_engine_specs.base import ( @@ -44,13 +46,63 @@ from superset.db_engine_specs.base import ( ) from superset.db_engine_specs.postgres import PostgresBaseEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import OAuth2TokenRefreshError from superset.models.sql_lab import Query +from superset.superset_typing import ( + OAuth2ClientConfig, + OAuth2State, +) from superset.utils import json from superset.utils.core import get_user_agent, QuerySource +from superset.utils.oauth2 import encode_oauth2_state, generate_code_challenge if TYPE_CHECKING: from superset.models.core import Database +try: + from snowflake.connector.errors import DatabaseError +except ImportError: + # Use a distinct sentinel type when snowflake is not installed to avoid + # matching unrelated exception types (using `Exception` would be too broad). + class _SnowflakeDatabaseError(Exception): + """Sentinel type to stand in for snowflake.connector.errors.DatabaseError.""" + + pass + + DatabaseError = _SnowflakeDatabaseError + + +class CustomSnowflakeAuthErrorMeta(type): + """ + Metaclass whose ``__instancecheck__`` matches Snowflake's invalid/expired + OAuth access-token error, so ``CustomSnowflakeAuthError`` can be used as the + ``oauth2_exception`` that triggers the OAuth2 re-auth dance. + + This is only honored via ``isinstance()`` (the path used by + ``BaseEngineSpec.needs_oauth2()``); ``except`` clauses do not call + ``__instancecheck__``, so it must not be relied on for exception catching. + """ + + def __instancecheck__(cls, instance: object) -> bool: + """ + Match Snowflake's invalid/expired OAuth token error, whether it arrives + wrapped by SQLAlchemy (e.g. ``Engine``-based execution) or as the raw + DBAPI exception — ``BaseEngineSpec.execute()`` runs against a bare + cursor and never wraps it, so both shapes must be handled here. + """ + orig: object = instance + if isinstance(instance, SqlalchemyDatabaseError): + orig = cast(SqlalchemyDatabaseError, instance).orig + + return isinstance(orig, DatabaseError) and "Invalid OAuth access token" in str( + orig + ) + + +class CustomSnowflakeAuthError(DatabaseError, metaclass=CustomSnowflakeAuthErrorMeta): + """Snowflake OAuth error type matched via the metaclass above (see note there).""" + + # Regular expressions to catch custom errors OBJECT_DOES_NOT_EXIST_REGEX = re.compile( r"Object (?P.*?) does not exist or not authorized." @@ -160,6 +212,7 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec): encrypted_extra_sensitive_fields = { "$.auth_params.privatekey_body": "Private Key Body", "$.auth_params.privatekey_pass": "Private Key Password", + "$.oauth2_client_info.secret": "OAuth2 Client Secret", } _time_grain_expressions = { @@ -198,6 +251,126 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec): ), } + # OAuth 2.0 support + supports_oauth2: bool = True + # `CustomSnowflakeAuthError` is only matched via `isinstance()` (see the + # metaclass docstring above), so it's paired with `OAuth2TokenRefreshError` + # (a real subclass) to keep `refresh_oauth2_token`'s `except` clause working. + oauth2_exception: type[Exception] | tuple[type[Exception], ...] = ( + CustomSnowflakeAuthError, + OAuth2TokenRefreshError, + ) + + @classmethod + def is_oauth2_enabled(cls) -> bool: + """ + Return whether OAuth2 authentication is enabled. + """ + + # When alerts or reports connect to the database in the background, + # OAuth2 authentication fails; therefore, OAuth2 authentication is disabled + # for background execution. + if not has_request_context(): + return False + + return ( + cls.supports_oauth2 + and cls.engine_name in app.config["DATABASE_OAUTH2_CLIENTS"] + ) + + @classmethod + def get_oauth2_config(cls) -> OAuth2ClientConfig | None: + """ + Build the DB engine spec level OAuth2 client config. + """ + if not cls.is_oauth2_enabled(): + return None + + return super().get_oauth2_config() + + @classmethod + def impersonate_user( + cls, + database: Database, + username: str | None, + user_token: str | None, + url: URL, + engine_kwargs: dict[str, Any], + ) -> tuple[URL, dict[str, Any]]: + """ + Modify URL and/or engine kwargs to impersonate a different user. + """ + connect_args: dict[str, Any] = engine_kwargs.setdefault("connect_args", {}) + + # When test_connection is executed (i.e., when validate_default_parameters is + # set to True in connect_args), authentication via OAuth is not performed. + # + # ``database.is_oauth2_enabled()`` returns True for a database-level OAuth2 + # client (``encrypted_extra.oauth2_client_info``) regardless of request + # context, unlike the app-config-based check in ``is_oauth2_enabled()`` + # above. Background executions (alerts/reports) have no per-user token, so + # ``has_request_context()`` must be checked explicitly here too, or OAuth + # gets switched on with no token to send. + if ( + not connect_args.get("validate_default_parameters", False) + and has_request_context() + and database.is_oauth2_enabled() + ): + url = url.update_query_dict({"authenticator": "oauth"}) + connect_args["authenticator"] = "oauth" + + if user_token: + if username is not None: + if is_feature_enabled("IMPERSONATE_WITH_EMAIL_PREFIX"): + # ``Database._get_sqla_engine()`` has already looked + # up the login and substituted the email prefix into + # ``username`` before calling this method when this + # flag is on. Looking it up again here as if it were + # still the login would fail whenever the two differ, + # leaving the default/service-account username paired + # with this user's OAuth token. Use it as given. + url = url.set(username=username) + else: + user = security_manager.find_user(username=username) + if user and user.email: + url = url.set(username=user.email) + + url = url.update_query_dict({"token": user_token}) + + return url, engine_kwargs + + @classmethod + def get_oauth2_authorization_uri( + cls, + config: OAuth2ClientConfig, + state: OAuth2State, + code_verifier: str | None = None, # pylint: disable=unused-argument + ) -> str: + """ + Return URI for initial OAuth2 request. + """ + uri = config["authorization_request_uri"] + # When calling the Snowflake OAuth authorization endpoint for a custom client, + # specify only the query parameters documented in the URL below. + # Adding unsupported parameters + # (e.g., `prompt` as used in BaseEngineSpec.get_oauth2_authorization_uri) + # will cause an error. + # https://docs.snowflake.com/user-guide/oauth-custom#query-parameters + params: dict[str, str] = { + "scope": config["scope"], + "response_type": "code", + "state": encode_oauth2_state(state), + "redirect_uri": config["redirect_uri"], + "client_id": config["id"], + } + + # Add PKCE parameters (RFC 7636) if code_verifier is provided + if code_verifier: + params["code_challenge"] = generate_code_challenge(code_verifier) + params["code_challenge_method"] = "S256" + + return parse.urljoin(uri, "?" + parse.urlencode(params)) + @staticmethod def get_extra_params( database: Database, source: QuerySource | None = None @@ -448,6 +621,18 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec): database: "Database", params: dict[str, Any], ) -> None: + # To use OAuth authentication, a database connection must first be created using + # another authenticator (typically key-pair authentication) + # with “Impersonate logged in user” enabled. + # Key-pair authentication is used for connection tests, + # while OAuth authentication is used when executing actual queries, + # such as in SQL Lab or dashboards. + # Therefore, when using OAuth authentication, the key-pair authentication + # settings are not loaded, and the connection is established using OAuth only. + connect_args: dict[str, Any] = params.get("connect_args") or {} + if connect_args.get("authenticator") == "oauth": + return + if not database.encrypted_extra: return try: diff --git a/tests/unit_tests/db_engine_specs/test_bigquery.py b/tests/unit_tests/db_engine_specs/test_bigquery.py index a1afed11b98..a35eaa0a893 100644 --- a/tests/unit_tests/db_engine_specs/test_bigquery.py +++ b/tests/unit_tests/db_engine_specs/test_bigquery.py @@ -715,12 +715,33 @@ def _patch_bq_fetch_deps( mocker: MockerFixture, max_mb: int = 200 ) -> tuple[mock.MagicMock, mock.MagicMock]: """Helper to patch Flask g and current_app for BigQuery fetch_data tests.""" - flask_g = mocker.patch("superset.db_engine_specs.bigquery.g") - app = mocker.patch("superset.db_engine_specs.bigquery.current_app") + # `new_callable=mock.MagicMock` is pinned explicitly rather than relying on + # ``mocker.patch``'s auto-detection of the mock class. That detection + # inspects whatever object currently sits at the patched attribute, so if + # an earlier test in the same session ever leaves an ``AsyncMock`` there + # (e.g. an improperly torn-down patch), every subsequent patch of the same + # attribute -- even ones created fresh here -- would also become an + # ``AsyncMock``, since ``AsyncMock`` classifies its own non-dunder child + # attributes as ``AsyncMock`` too. Pinning the callable sidesteps that + # self-perpetuating class inference entirely. + flask_g = mocker.patch( + "superset.db_engine_specs.bigquery.g", new_callable=mock.MagicMock + ) + app = mocker.patch( + "superset.db_engine_specs.bigquery.current_app", new_callable=mock.MagicMock + ) # Make current_app truthy and .config.get() return a plain int app.__bool__ = mock.Mock(return_value=True) app.config = mock.MagicMock() app.config.get = mock.Mock(return_value=max_mb) + # ``fetch_data`` only records ``g.bq_memory_limited*`` when + # ``has_request_context()`` is true. Outside of a real Flask request + # (as in these unit tests) that's always false, so without patching it + # the assignments never happen and the mocked ``g`` attributes stay + # unset child mocks instead of the expected booleans/counts. + mocker.patch( + "superset.db_engine_specs.bigquery.has_request_context", return_value=True + ) return flask_g, app diff --git a/tests/unit_tests/db_engine_specs/test_snowflake.py b/tests/unit_tests/db_engine_specs/test_snowflake.py index 3f166b22650..d41c8cbc4b7 100644 --- a/tests/unit_tests/db_engine_specs/test_snowflake.py +++ b/tests/unit_tests/db_engine_specs/test_snowflake.py @@ -23,9 +23,11 @@ from unittest import mock import pytest from pytest_mock import MockerFixture -from sqlalchemy.engine.url import make_url +from sqlalchemy.engine.url import make_url, URL +from superset.app import SupersetApp from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.superset_typing import OAuth2ClientConfig from superset.utils import json from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm from tests.unit_tests.fixtures.common import dttm # noqa: F401 @@ -350,6 +352,30 @@ def test_mask_encrypted_extra() -> None: ) +def test_mask_encrypted_extra_oauth2_client_secret() -> None: + """ + The database-level OAuth2 client secret must be masked in + ``masked_encrypted_extra``, matching the other engine specs supporting + the same ``oauth2_client_info`` path (gsheets, trino) -- otherwise a + database editor can read it back unmasked. + """ + from superset.db_engine_specs.snowflake import SnowflakeEngineSpec + + config = json.dumps( + { + "auth_method": "oauth2", + "oauth2_client_info": {"id": "client-id", "secret": "my-secret"}, + } + ) + + assert SnowflakeEngineSpec.mask_encrypted_extra(config) == json.dumps( + { + "auth_method": "oauth2", + "oauth2_client_info": {"id": "client-id", "secret": "XXXXXXXXXX"}, + } + ) + + def test_mask_encrypted_extra_no_fields() -> None: """ Test that the private key is masked when the database is edited. @@ -462,3 +488,278 @@ def test_unmask_encrypted_extra() -> None: }, } ) + + +@pytest.fixture +def oauth2_config() -> OAuth2ClientConfig: + """ + Config for Snowflake OAuth2. + """ + return { + "id": "snowflake-oauth2-client-id", + "secret": "snowflake-oauth2-client-secret", + "scope": "refresh_token", + "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/", + "authorization_request_uri": "https://snowflake.oauth2.example/oauth/authorize", + "token_request_uri": "https://snowflake.oauth2.example/oauth/token-request", + "request_content_type": "data", + } + + +def test_get_oauth2_token( + mocker: MockerFixture, + oauth2_config: OAuth2ClientConfig, +) -> None: + """ + Test `get_oauth2_token`. + """ + from superset.db_engine_specs.snowflake import SnowflakeEngineSpec + + requests: mock.MagicMock = mocker.patch("superset.db_engine_specs.base.requests") + requests.post().json.return_value = { + "access_token": "access-token", + "expires_in": 3600, + "scope": "scope", + "token_type": "Bearer", + "refresh_token": "refresh-token", + } + + assert SnowflakeEngineSpec.get_oauth2_token(oauth2_config, "code") == { + "access_token": "access-token", + "expires_in": 3600, + "scope": "scope", + "token_type": "Bearer", + "refresh_token": "refresh-token", + } + requests.post.assert_called_with( + "https://snowflake.oauth2.example/oauth/token-request", + data={ + "code": "code", + "client_id": "snowflake-oauth2-client-id", + "client_secret": "snowflake-oauth2-client-secret", + "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/", + "grant_type": "authorization_code", + }, + timeout=30.0, + ) + + +def test_impersonate_user(app: SupersetApp, mocker: MockerFixture) -> None: + """ + Test that Snowflake supports user impersonation. + + Impersonation only applies within a request context (see + ``test_impersonate_user_outside_request_context`` below for the + background-execution case), so these assertions run inside one. + """ + from superset.db_engine_specs.snowflake import SnowflakeEngineSpec + from superset.models.core import Database + + database: Database = Database(sqlalchemy_uri="snowflake://abc") + + mocker.patch( + "superset.db_engine_specs.snowflake.SnowflakeEngineSpec.is_oauth2_enabled", + return_value=True, + ) + + with app.test_request_context("/some/place/"): + assert SnowflakeEngineSpec.impersonate_user( + database=database, + username=None, + user_token=None, + url=make_url("snowflake://user:pass@account/database_name/default"), + engine_kwargs={ + "connect_args": { + "validate_default_parameters": True, + }, + }, + ) == ( + make_url("snowflake://user:pass@account/database_name/default"), + {"connect_args": {"validate_default_parameters": True}}, + ) + + assert SnowflakeEngineSpec.impersonate_user( + database=database, + username=None, + user_token=None, + url=make_url("snowflake://user:pass@account/database_name/default"), + engine_kwargs={}, + ) == ( + make_url( + "snowflake://user:pass@account/database_name/default?authenticator=oauth" + ), + {"connect_args": {"authenticator": "oauth"}}, + ) + + mocker.patch( + "superset.db_engine_specs.snowflake.is_feature_enabled", + return_value=True, + ) + + mocker.patch( + "superset.security_manager.find_user", + return_value=mocker.MagicMock(email="impersonated_user@example.com"), + ) + assert SnowflakeEngineSpec.impersonate_user( + database=database, + username="impersonated_user", + user_token="test_token", # noqa: S106 + url=make_url("snowflake://user:pass@account/database_name/default"), + engine_kwargs={}, + ) == ( + make_url( + "snowflake://impersonated_user:pass@account/database_name/default?authenticator=oauth&token=test_token" + ), + {"connect_args": {"authenticator": "oauth"}}, + ) + + +def test_impersonate_user_email_prefix_uses_username_directly( + app: SupersetApp, mocker: MockerFixture +) -> None: + """ + With IMPERSONATE_WITH_EMAIL_PREFIX enabled, ``Database._get_sqla_engine()`` + has already substituted the email prefix for the login username before + calling ``impersonate_user`` -- the value it passes in is no longer a + lookupable login. Re-looking it up as a username (the pre-fix behavior) + fails whenever the login differs from the prefix, silently leaving the + default/service-account username paired with the impersonated user's + OAuth token instead of failing loudly. The fixed code must use the given + value directly and must not call ``find_user`` at all in this branch. + """ + from superset.db_engine_specs.snowflake import SnowflakeEngineSpec + from superset.models.core import Database + + database: Database = Database(sqlalchemy_uri="snowflake://abc") + + mocker.patch( + "superset.db_engine_specs.snowflake.SnowflakeEngineSpec.is_oauth2_enabled", + return_value=True, + ) + mocker.patch( + "superset.db_engine_specs.snowflake.is_feature_enabled", + return_value=True, + ) + find_user = mocker.patch("superset.security_manager.find_user") + + with app.test_request_context("/some/place/"): + # "jdoe" is the email prefix Database._get_sqla_engine() already + # derived; the login it derived it from ("jdoe123", say) is gone by + # this point and must not be re-derived here. + result = SnowflakeEngineSpec.impersonate_user( + database=database, + username="jdoe", + user_token="test_token", # noqa: S106 + url=make_url("snowflake://user:pass@account/database_name/default"), + engine_kwargs={}, + ) + + assert result == ( + make_url( + "snowflake://jdoe:pass@account/database_name/default?authenticator=oauth&token=test_token" + ), + {"connect_args": {"authenticator": "oauth"}}, + ) + find_user.assert_not_called() + + +def test_impersonate_user_outside_request_context(mocker: MockerFixture) -> None: + """ + Background executions (alerts/reports) have no per-user token, so OAuth + impersonation must not engage outside a request context — even when + ``database.is_oauth2_enabled()`` returns True because of a + database-level OAuth2 client config, which (unlike the app-config-based + check) isn't itself request-context-aware. + """ + from superset.db_engine_specs.snowflake import SnowflakeEngineSpec + from superset.models.core import Database + + database: Database = Database(sqlalchemy_uri="snowflake://abc") + mocker.patch.object(Database, "is_oauth2_enabled", return_value=True) + + url: URL = make_url("snowflake://user:pass@account/database_name/default") + assert SnowflakeEngineSpec.impersonate_user( + database=database, + username=None, + user_token="test_token", # noqa: S106 + url=url, + engine_kwargs={}, + ) == (url, {"connect_args": {}}) + + +def test_custom_snowflake_auth_error_matches_raw_dbapi_exception() -> None: + """ + `BaseEngineSpec.execute()` runs against a bare DBAPI cursor, so the + exception it sees is the raw Snowflake error, never wrapped by + SQLAlchemy. `CustomSnowflakeAuthError` must still recognize it so the + OAuth2 re-auth dance triggers for SQL Lab queries. + """ + from superset.db_engine_specs.snowflake import ( + CustomSnowflakeAuthError, + DatabaseError, + ) + + raw_error: Exception = DatabaseError("250001: Invalid OAuth access token.") + assert isinstance(raw_error, CustomSnowflakeAuthError) + + +def test_custom_snowflake_auth_error_matches_sqlalchemy_wrapped_exception() -> None: + """ + Some call sites execute through SQLAlchemy's `Engine`, which wraps the + original DBAPI exception in `sqlalchemy.exc.DatabaseError.orig`. + `CustomSnowflakeAuthError` must keep matching this shape too. + """ + from sqlalchemy.exc import DatabaseError as SqlalchemyDatabaseError + + from superset.db_engine_specs.snowflake import ( + CustomSnowflakeAuthError, + DatabaseError, + ) + + wrapped_error: SqlalchemyDatabaseError = SqlalchemyDatabaseError( + statement="SELECT 1", + params=None, + orig=DatabaseError("250001: Invalid OAuth access token."), + ) + assert isinstance(wrapped_error, CustomSnowflakeAuthError) + + +def test_custom_snowflake_auth_error_does_not_match_unrelated_errors() -> None: + """ + Other Snowflake DB errors, and non-Snowflake exceptions, must not be + mistaken for an expired OAuth token. + """ + from superset.db_engine_specs.snowflake import ( + CustomSnowflakeAuthError, + DatabaseError, + ) + + assert not isinstance( + DatabaseError("Object FOO does not exist."), CustomSnowflakeAuthError + ) + assert not isinstance( + ValueError("Invalid OAuth access token."), CustomSnowflakeAuthError + ) + + +def test_snowflake_oauth2_exception_catches_refresh_token_error() -> None: + """ + `refresh_oauth2_token()` catches failures from the (unoverridden) base + `get_oauth2_fresh_token()` with `except db_engine_spec.oauth2_exception`. + That base method raises `OAuth2TokenRefreshError`, which isn't related to + `CustomSnowflakeAuthError` by real subclassing, so `oauth2_exception` must + include it directly -- an `except` clause never triggers the metaclass's + `__instancecheck__`, unlike `isinstance()`. + """ + from superset.db_engine_specs.snowflake import SnowflakeEngineSpec + from superset.exceptions import OAuth2TokenRefreshError + + try: + raise OAuth2TokenRefreshError("refresh token revoked") + except SnowflakeEngineSpec.oauth2_exception: + pass + else: + pytest.fail( + "OAuth2TokenRefreshError must be caught by " + "SnowflakeEngineSpec.oauth2_exception" + ) diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py index bb9a2c1cbb8..008a81441c8 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/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: