You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@superset.apache.org by "sadpandajoe (via GitHub)" <gi...@apache.org> on 2026/08/21 07:16:46 UTC
Re: [PR] feat: upstream oauth token forwarding [superset]
sadpandajoe commented on code in PR #38469:
URL: https://github.com/apache/superset/pull/38469#discussion_r3828166357
##########
superset/utils/oauth2.py:
##########
@@ -318,3 +318,186 @@ def check_for_oauth2(database: Database) -> Iterator[None]:
if database.is_oauth2_enabled() and database.db_engine_spec.needs_oauth2(ex):
database.db_engine_spec.start_oauth2_dance(database)
raise
+
+
+def get_access_token_for_database(database: Database, user_id: int) -> str | None:
+ """
+ Return a valid OAuth2 access token for the given database and user.
+
+ First checks if the database has an upstream OAuth provider configured in its
+ ``extra`` JSON (key ``oauth2_upstream_provider``). If so, returns the saved
+ upstream login token for that provider.
+
+ Otherwise, falls back to the database-specific OAuth2 flow.
+ """
+ upstream_provider = database.get_extra().get("oauth2_upstream_provider")
+ if upstream_provider:
+ access_token = get_upstream_provider_token(upstream_provider, user_id)
+ if access_token:
+ logger.info(
+ "Using upstream OAuth token from provider '%s' "
+ "for database '%s' (id=%s, user_id=%d)",
+ upstream_provider,
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return access_token
+ logger.warning(
+ "Upstream provider '%s' configured for database '%s' "
+ "(id=%s) but no valid token found for user_id=%d, "
+ "falling back to database-specific OAuth2",
+ upstream_provider,
+ database.database_name,
+ database.id,
+ user_id,
+ )
+
+ # Fall back to database-specific OAuth2 (also used when upstream token is
+ # unavailable, e.g. expired without a refresh token).
+ # Database-specific OAuth2 requires a persisted database (needs database.id
+ # to look up stored tokens).
+ if database.id is None:
+ logger.debug(
+ "Database '%s' has no persisted id, "
+ "skipping database-specific OAuth2 token lookup",
+ database.database_name,
+ )
+ return None
+
+ oauth2_config = database.get_oauth2_config()
+ if oauth2_config:
+ logger.info(
+ "Using database-specific OAuth2 token "
+ "for database '%s' (id=%d, user_id=%d)",
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return get_oauth2_access_token(
+ oauth2_config, database.id, user_id, database.db_engine_spec
+ )
+
+ logger.warning(
+ "No OAuth2 token available for database '%s' (id=%d, user_id=%d)",
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return None
+
+
+def save_user_provider_token(
+ user_id: int,
+ provider: str,
+ token_response: dict[str, Any],
+) -> None:
+ """
+ Upsert an UpstreamOAuthToken row for the given user + provider.
+ """
+ from superset.models.core import UpstreamOAuthToken
+
+ token: UpstreamOAuthToken | None = (
+ db.session.query(UpstreamOAuthToken)
+ .filter_by(user_id=user_id, provider=provider)
+ .one_or_none()
+ )
+ if token is None:
+ token = UpstreamOAuthToken(user_id=user_id, provider=provider)
+
+ token.access_token = token_response.get("access_token")
Review Comment:
FAB supports a provider-specific `token_key`, but this persistence path always reads `access_token`. Providers configured with another token field can log in successfully while this stores `None`, so upstream forwarding never receives their token. Could this resolve the configured token key before persisting it?
##########
superset/utils/oauth2.py:
##########
@@ -318,3 +318,186 @@ def check_for_oauth2(database: Database) -> Iterator[None]:
if database.is_oauth2_enabled() and database.db_engine_spec.needs_oauth2(ex):
database.db_engine_spec.start_oauth2_dance(database)
raise
+
+
+def get_access_token_for_database(database: Database, user_id: int) -> str | None:
+ """
+ Return a valid OAuth2 access token for the given database and user.
+
+ First checks if the database has an upstream OAuth provider configured in its
+ ``extra`` JSON (key ``oauth2_upstream_provider``). If so, returns the saved
+ upstream login token for that provider.
+
+ Otherwise, falls back to the database-specific OAuth2 flow.
+ """
+ upstream_provider = database.get_extra().get("oauth2_upstream_provider")
+ if upstream_provider:
+ access_token = get_upstream_provider_token(upstream_provider, user_id)
+ if access_token:
+ logger.info(
+ "Using upstream OAuth token from provider '%s' "
+ "for database '%s' (id=%s, user_id=%d)",
+ upstream_provider,
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return access_token
+ logger.warning(
+ "Upstream provider '%s' configured for database '%s' "
+ "(id=%s) but no valid token found for user_id=%d, "
+ "falling back to database-specific OAuth2",
+ upstream_provider,
+ database.database_name,
+ database.id,
+ user_id,
+ )
+
+ # Fall back to database-specific OAuth2 (also used when upstream token is
+ # unavailable, e.g. expired without a refresh token).
+ # Database-specific OAuth2 requires a persisted database (needs database.id
+ # to look up stored tokens).
+ if database.id is None:
+ logger.debug(
+ "Database '%s' has no persisted id, "
+ "skipping database-specific OAuth2 token lookup",
+ database.database_name,
+ )
+ return None
+
+ oauth2_config = database.get_oauth2_config()
+ if oauth2_config:
+ logger.info(
+ "Using database-specific OAuth2 token "
+ "for database '%s' (id=%d, user_id=%d)",
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return get_oauth2_access_token(
+ oauth2_config, database.id, user_id, database.db_engine_spec
+ )
+
+ logger.warning(
Review Comment:
`_get_sqla_engine()` now calls this helper for every authenticated engine creation, including databases that do not use OAuth. That normal path reaches this warning on every query/chart render, flooding production logs with a false failure signal. Should the no-OAuth case return quietly (and reserve warnings for configured OAuth failures)?
##########
superset/utils/oauth2.py:
##########
@@ -318,3 +318,186 @@ def check_for_oauth2(database: Database) -> Iterator[None]:
if database.is_oauth2_enabled() and database.db_engine_spec.needs_oauth2(ex):
database.db_engine_spec.start_oauth2_dance(database)
raise
+
+
+def get_access_token_for_database(database: Database, user_id: int) -> str | None:
+ """
+ Return a valid OAuth2 access token for the given database and user.
+
+ First checks if the database has an upstream OAuth provider configured in its
+ ``extra`` JSON (key ``oauth2_upstream_provider``). If so, returns the saved
+ upstream login token for that provider.
+
+ Otherwise, falls back to the database-specific OAuth2 flow.
+ """
+ upstream_provider = database.get_extra().get("oauth2_upstream_provider")
+ if upstream_provider:
+ access_token = get_upstream_provider_token(upstream_provider, user_id)
+ if access_token:
+ logger.info(
+ "Using upstream OAuth token from provider '%s' "
+ "for database '%s' (id=%s, user_id=%d)",
+ upstream_provider,
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return access_token
+ logger.warning(
+ "Upstream provider '%s' configured for database '%s' "
+ "(id=%s) but no valid token found for user_id=%d, "
+ "falling back to database-specific OAuth2",
+ upstream_provider,
+ database.database_name,
+ database.id,
+ user_id,
+ )
+
+ # Fall back to database-specific OAuth2 (also used when upstream token is
+ # unavailable, e.g. expired without a refresh token).
+ # Database-specific OAuth2 requires a persisted database (needs database.id
+ # to look up stored tokens).
+ if database.id is None:
+ logger.debug(
+ "Database '%s' has no persisted id, "
+ "skipping database-specific OAuth2 token lookup",
+ database.database_name,
+ )
+ return None
+
+ oauth2_config = database.get_oauth2_config()
+ if oauth2_config:
+ logger.info(
+ "Using database-specific OAuth2 token "
+ "for database '%s' (id=%d, user_id=%d)",
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return get_oauth2_access_token(
+ oauth2_config, database.id, user_id, database.db_engine_spec
+ )
+
+ logger.warning(
+ "No OAuth2 token available for database '%s' (id=%d, user_id=%d)",
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return None
+
+
+def save_user_provider_token(
+ user_id: int,
+ provider: str,
+ token_response: dict[str, Any],
+) -> None:
+ """
+ Upsert an UpstreamOAuthToken row for the given user + provider.
+ """
+ from superset.models.core import UpstreamOAuthToken
+
+ token: UpstreamOAuthToken | None = (
+ db.session.query(UpstreamOAuthToken)
+ .filter_by(user_id=user_id, provider=provider)
+ .one_or_none()
+ )
+ if token is None:
+ token = UpstreamOAuthToken(user_id=user_id, provider=provider)
+
+ token.access_token = token_response.get("access_token")
+ expires_in = token_response.get("expires_in")
+ token.access_token_expiration = (
+ datetime.now() + timedelta(seconds=expires_in) if expires_in else None
+ )
+ token.refresh_token = token_response.get("refresh_token")
+ db.session.add(token)
+ db.session.commit()
+
+
+@backoff.on_exception(
+ backoff.expo,
+ AcquireDistributedLockFailedException,
+ factor=10,
+ base=2,
+ max_tries=5,
+)
+def get_upstream_provider_token(provider: str, user_id: int) -> str | None:
+ """
+ Retrieve a valid access token for the given provider and user.
+
+ If the token is expired and a refresh token exists, attempt to refresh it.
+ Returns None if no valid token is available.
+ """
+ from superset.models.core import UpstreamOAuthToken
+
+ token: UpstreamOAuthToken | None = (
+ db.session.query(UpstreamOAuthToken)
+ .filter_by(user_id=user_id, provider=provider)
+ .one_or_none()
+ )
+ if token is None:
+ return None
+
+ now = datetime.now()
+ if token.access_token_expiration is None or token.access_token_expiration > now:
+ return token.access_token
+
+ # Token is expired
+ if token.refresh_token:
+ return _refresh_upstream_provider_token(token, provider)
+
+ db.session.delete(token)
+ db.session.commit()
+ return None
+
+
+def _refresh_upstream_provider_token(
+ token: UpstreamOAuthToken,
+ provider: str,
+) -> str | None:
+ """
+ Use the refresh token to obtain a new access token from the provider.
+ Updates and persists the token if successful; deletes it on failure.
+ """
+ from flask import current_app as flask_app
+
+ with DistributedLock(
+ namespace="refresh_upstream_oauth_token",
+ user_id=token.user_id,
+ provider=provider,
+ ):
+ try:
+ remote_app = flask_app.extensions["authlib.integrations.flask_client"][
Review Comment:
This extension entry is Authlib's OAuth registry, not a mapping of configured remotes, so subscripting it raises before `fetch_access_token()` is reached. The broad handler then deletes the stored token; every expired upstream token with a refresh token is therefore lost. Should this use the configured OAuth remote accessor instead?
##########
superset/utils/oauth2.py:
##########
@@ -318,3 +318,186 @@ def check_for_oauth2(database: Database) -> Iterator[None]:
if database.is_oauth2_enabled() and database.db_engine_spec.needs_oauth2(ex):
database.db_engine_spec.start_oauth2_dance(database)
raise
+
+
+def get_access_token_for_database(database: Database, user_id: int) -> str | None:
+ """
+ Return a valid OAuth2 access token for the given database and user.
+
+ First checks if the database has an upstream OAuth provider configured in its
+ ``extra`` JSON (key ``oauth2_upstream_provider``). If so, returns the saved
+ upstream login token for that provider.
+
+ Otherwise, falls back to the database-specific OAuth2 flow.
+ """
+ upstream_provider = database.get_extra().get("oauth2_upstream_provider")
+ if upstream_provider:
+ access_token = get_upstream_provider_token(upstream_provider, user_id)
+ if access_token:
+ logger.info(
+ "Using upstream OAuth token from provider '%s' "
+ "for database '%s' (id=%s, user_id=%d)",
+ upstream_provider,
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return access_token
+ logger.warning(
+ "Upstream provider '%s' configured for database '%s' "
+ "(id=%s) but no valid token found for user_id=%d, "
+ "falling back to database-specific OAuth2",
+ upstream_provider,
+ database.database_name,
+ database.id,
+ user_id,
+ )
+
+ # Fall back to database-specific OAuth2 (also used when upstream token is
+ # unavailable, e.g. expired without a refresh token).
+ # Database-specific OAuth2 requires a persisted database (needs database.id
+ # to look up stored tokens).
+ if database.id is None:
+ logger.debug(
+ "Database '%s' has no persisted id, "
+ "skipping database-specific OAuth2 token lookup",
+ database.database_name,
+ )
+ return None
+
+ oauth2_config = database.get_oauth2_config()
+ if oauth2_config:
+ logger.info(
+ "Using database-specific OAuth2 token "
+ "for database '%s' (id=%d, user_id=%d)",
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return get_oauth2_access_token(
+ oauth2_config, database.id, user_id, database.db_engine_spec
+ )
+
+ logger.warning(
+ "No OAuth2 token available for database '%s' (id=%d, user_id=%d)",
+ database.database_name,
+ database.id,
+ user_id,
+ )
+ return None
+
+
+def save_user_provider_token(
+ user_id: int,
+ provider: str,
+ token_response: dict[str, Any],
+) -> None:
+ """
+ Upsert an UpstreamOAuthToken row for the given user + provider.
+ """
+ from superset.models.core import UpstreamOAuthToken
+
+ token: UpstreamOAuthToken | None = (
+ db.session.query(UpstreamOAuthToken)
+ .filter_by(user_id=user_id, provider=provider)
+ .one_or_none()
+ )
+ if token is None:
+ token = UpstreamOAuthToken(user_id=user_id, provider=provider)
+
+ token.access_token = token_response.get("access_token")
+ expires_in = token_response.get("expires_in")
+ token.access_token_expiration = (
+ datetime.now() + timedelta(seconds=expires_in) if expires_in else None
+ )
+ token.refresh_token = token_response.get("refresh_token")
Review Comment:
A second authorization-code login is allowed to return a new access token without a `refresh_token`. This assignment then erases the still-valid stored refresh token, so the next expiry cannot be refreshed and the forwarding flow falls back to another OAuth dance. Should this preserve the existing refresh token unless the response explicitly supplies a replacement?
--
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: notifications-unsubscribe@superset.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org