rusackas commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3846884066
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -192,6 +244,119 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
),
}
+ # OAuth 2.0 support
Review Comment:
Good catch, fixed. Registered `$.oauth2_client_info.secret` in
`encrypted_extra_sensitive_fields`, matching gsheets/trino. Added a regression
test asserting the secret is masked in `masked_encrypted_extra`.
##########
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:
Good catch, fixed. `impersonate_user` no longer calls `find_user` on the
value passed in when `IMPERSONATE_WITH_EMAIL_PREFIX` is enabled -- it's already
the email prefix by the time it gets here, so re-deriving it as a login was
silently failing whenever the two differ. Uses the given value directly in that
case now; kept the existing login->email lookup for the non-prefix path, where
the raw login is what's actually passed in. Added a regression test covering
the prefix path and asserting `find_user` isn't called there.
--
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]
