rusackas commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3564845976
##########
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():
Review Comment:
Already addressed (duplicate of the same finding on an earlier push) — see
`impersonate_user`'s `database.is_oauth2_enabled()` gate.
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -191,6 +224,99 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
),
}
+ # OAuth 2.0 support
+ supports_oauth2 = True
Review Comment:
Not applying — every other boolean class attribute in `SnowflakeEngineSpec`
(`supports_dynamic_schema`, `force_column_alias_quotes`,
`use_equality_for_boolean_filters`, etc.) is unannotated too; singling out
`supports_oauth2` would be inconsistent with the surrounding style, and mypy
already infers it fine.
##########
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", {})
Review Comment:
Applied — `connect_args: dict[str, Any] = engine_kwargs.setdefault(...)`
(08d490e29a), matching the same pattern already used in `get_extra_params` in
this file.
##########
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"],
+ }
Review Comment:
Applied — `params: dict[str, str] = {...}` (08d490e29a); all of
`OAuth2ClientConfig`'s relevant fields are `str`, so this is accurate.
--
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]
