rusackas commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3564846263
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -191,6 +224,99 @@ 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 = 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 cls.is_oauth2_enabled()
+ ):
+ url = url.update_query_dict({"authenticator": "oauth"})
+ connect_args["authenticator"] = "oauth"
+
+ if user_token and cls.is_oauth2_enabled():
+ if username is not None:
+ user = security_manager.find_user(username=username)
+ if user and user.email:
+ if is_feature_enabled("IMPERSONATE_WITH_EMAIL_PREFIX"):
+ url = url.set(username=user.email.split("@")[0])
+ else:
+ url = url.set(username=user.email)
+
+ url = url.update_query_dict({"token": user_token})
+
Review Comment:
Already addressed — the whole OAuth block in `impersonate_user` (setting the
authenticator *and* injecting the token) is now gated behind `not
connect_args.get("validate_default_parameters", False)`, so test-connection
flows never get mixed auth params.
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -44,12 +46,43 @@
from superset.db_engine_specs.postgres import PostgresBaseEngineSpec
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
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
Review Comment:
Tried this — it actually breaks mypy: `DatabaseError: type[Exception] =
_SnowflakeDatabaseError` triggers `error: Name "DatabaseError" already defined
(possibly by an import) [no-redef]`, since the try branch already binds
`DatabaseError` via the conditional import. Not applying.
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -44,12 +46,53 @@
from superset.db_engine_specs.postgres import PostgresBaseEngineSpec
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
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:
+ return (
+ isinstance(instance, SqlalchemyDatabaseError)
+ and isinstance(cast(SqlalchemyDatabaseError, instance).orig,
DatabaseError)
+ and "Invalid OAuth access token" in str(instance)
+ )
Review Comment:
Confirmed real, and applied (08d490e29a). Verified the path: SQL Lab query
execution routes through `BaseEngineSpec.execute()`, which calls
`cursor.execute(query)` directly on a raw DBAPI cursor (from
`Database.get_raw_connection()`) — so the exception it catches is the unwrapped
snowflake-connector error, never a `sqlalchemy.exc.DatabaseError`.
`__instancecheck__` now checks `.orig` when the instance is SQLAlchemy-wrapped,
but falls back to checking the raw exception directly otherwise. Added
regression tests covering both shapes plus a negative case; the raw-DBAPI test
fails against the pre-fix code (verified by stashing the fix and re-running).
--
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]
