rusackas commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3564845633
##########
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
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):
+ 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:
Applied — `__instancecheck__` now has a docstring explaining the
raw-vs-SQLAlchemy-wrapped matching it does (08d490e29a), added alongside a real
fix to what it matches (see the reply on the newer critical-severity comment
below on this same method).
##########
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"]
+ )
Review Comment:
Already addressed — `impersonate_user` gates on
`database.is_oauth2_enabled()` (an instance method that checks per-database
`encrypted_extra` config first, falling back to the class-level check), not
just `cls.is_oauth2_enabled()`. Confirmed by reading
`Database.is_oauth2_enabled()` in `superset/models/core.py`, which does exactly
this.
##########
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})
+
+ 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 = {
+ "scope": config["scope"],
+ "response_type": "code",
+ "state": encode_oauth2_state(state),
+ "redirect_uri": config["redirect_uri"],
+ "client_id": config["id"],
+ }
+ return parse.urljoin(uri, "?" + parse.urlencode(params))
Review Comment:
Already addressed — `get_oauth2_authorization_uri` now adds
`code_challenge`/`code_challenge_method` when `code_verifier` is provided.
--
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]
