rusackas commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3564844477
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -135,6 +161,71 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
),
}
+ # OAuth 2.0 support
+ supports_oauth2 = True
+ oauth2_exception = CustomSnowflakeAuthError
+
+ @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:
+ 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,
+ ) -> 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:
Checked this — it's not actually a bug. `urljoin(base, "?query")` follows
RFC 3986 §5.3: a relative reference consisting only of a query component
resolves against the base URI by keeping its path and replacing only the query.
Verified in Python directly:
`urljoin('https://acct.snowflakecomputing.com/oauth/authorize',
'?scope=s&client_id=c')` →
`https://acct.snowflakecomputing.com/oauth/authorize?scope=s&client_id=c` —
path preserved, query replaced, exactly as intended. Not applying the suggested
rewrite.
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -33,19 +33,45 @@
from sqlalchemy import types
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.engine.url import URL
+from sqlalchemy.exc import DatabaseError as SqlalchemyDatabaseError
+from superset import is_feature_enabled, security_manager
from superset.constants import TimeGrain
from superset.databases.utils import make_url_safe
from superset.db_engine_specs.base import BaseEngineSpec, BasicPropertiesType
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:
+ DatabaseError = Exception
Review Comment:
Already addressed — `CustomSnowflakeAuthErrorMeta` has a class docstring
explaining its purpose and the `__instancecheck__`-only caveat.
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -377,6 +468,18 @@ def update_params_from_encrypted_extra(
database: "Database",
params: dict[str, Any],
) -> None:
+ # To use OAuth authentication, a database connection must first be
created using
+ # another authenticator (typically key-pair authentication)
+ # with “Impersonate logged in user” enabled.
+ # Key-pair authentication is used for connection tests,
+ # while OAuth authentication is used when executing actual queries,
+ # such as in SQL Lab or dashboards.
+ # Therefore, when using OAuth authentication, the key-pair
authentication
+ # settings are not loaded, and the connection is established using
OAuth only.
+ connects_args = params.get("connect_args", {})
+ if connects_args.get("authenticator") == "oauth":
Review Comment:
Already addressed — `CustomSnowflakeAuthError` has a one-line docstring
pointing to the metaclass for details.
--
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]
