rusackas commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3565173810
##########
tests/unit_tests/db_engine_specs/test_snowflake.py:
##########
@@ -438,3 +439,177 @@ def test_unmask_encrypted_extra() -> None:
},
}
)
+
+
[email protected]
+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 = 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(mocker: MockerFixture) -> None:
+ """
+ Test that Snowflake supports user impersonation.
+ """
+ from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
+ from superset.models.core import Database
+
+ database = Database(sqlalchemy_uri="snowflake://abc")
+
+ mocker.patch(
+
"superset.db_engine_specs.snowflake.SnowflakeEngineSpec.is_oauth2_enabled",
+ return_value=True,
+ )
+
+ 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="[email protected]"),
+ )
+ 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_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 = DatabaseError("250001: Invalid OAuth access token.")
Review Comment:
Applied — `raw_error: Exception = DatabaseError(...)` (8ba0c75b39).
##########
tests/unit_tests/db_engine_specs/test_snowflake.py:
##########
@@ -438,3 +439,177 @@ def test_unmask_encrypted_extra() -> None:
},
}
)
+
+
[email protected]
+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 = 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(mocker: MockerFixture) -> None:
+ """
+ Test that Snowflake supports user impersonation.
+ """
+ from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
+ from superset.models.core import Database
+
+ database = Database(sqlalchemy_uri="snowflake://abc")
+
+ mocker.patch(
+
"superset.db_engine_specs.snowflake.SnowflakeEngineSpec.is_oauth2_enabled",
+ return_value=True,
+ )
+
+ 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="[email protected]"),
+ )
+ 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_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 = 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(
+ statement="SELECT 1",
+ params=None,
+ orig=DatabaseError("250001: Invalid OAuth access token."),
+ )
Review Comment:
Applied — `wrapped_error: SqlalchemyDatabaseError =
SqlalchemyDatabaseError(...)` (8ba0c75b39).
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -191,6 +242,105 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
),
}
+ # OAuth 2.0 support
+ supports_oauth2 = True
+ oauth2_exception = CustomSnowflakeAuthError
+
+ @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.
+ if (
+ not connect_args.get("validate_default_parameters", False)
+ and database.is_oauth2_enabled()
+ ):
Review Comment:
Confirmed real, and applied (8ba0c75b39). `database.is_oauth2_enabled()`
returns True for a database-level OAuth2 client
(`encrypted_extra.oauth2_client_info`) regardless of request context — that
path bypasses the app-config-based `has_request_context()` guard in
`SnowflakeEngineSpec.is_oauth2_enabled()` entirely (this shared method lives in
`superset/models/core.py`, unmodified by this PR, so it's used by other
OAuth2-capable engines too). Added an explicit `has_request_context()` check
directly in `impersonate_user()` so the background-execution guarantee holds
for both config sources. Added a regression test
(`test_impersonate_user_outside_request_context`) confirming OAuth doesn't
engage outside a request context even when `is_oauth2_enabled()` is mocked to
return True — it fails against the pre-fix code (verified locally). Also had to
fix `test_impersonate_user` itself, which never actually ran inside a request
context before, so it never caught this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
