sadpandajoe commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3821761878
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -192,6 +244,119 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
),
}
+ # OAuth 2.0 support
Review Comment:
This enables Snowflake's database-level OAuth client path, but
`encrypted_extra_sensitive_fields` still omits `$.oauth2_client_info.secret`
while the other specs supporting that path mask it. A database editor can then
receive the configured client secret through `masked_encrypted_extra`. Could
this register that field as sensitive before enabling the path?
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -192,6 +244,119 @@ 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:
+ user = security_manager.find_user(username=username)
Review Comment:
With `IMPERSONATE_WITH_EMAIL_PREFIX` enabled, `Database._get_sqla_engine()`
has already converted the effective username to the email prefix before calling
this method. Looking that prefix up again as a username fails whenever the
login differs from the prefix, leaving the service-account username paired with
the user's OAuth token. Could this use the supplied username directly instead
of repeating the lookup/mapping?
--
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]
