sadpandajoe commented on code in PR #43623:
URL: https://github.com/apache/superset/pull/43623#discussion_r3898624706
##########
superset/utils/oauth2.py:
##########
@@ -202,11 +257,105 @@ def refresh_oauth2_token(
if new_refresh_token := token_response.get("refresh_token"):
token.refresh_token = new_refresh_token
- db.session.add(token)
+ token_session.add(token)
+ if force:
+ # Make rotated access and refresh tokens visible to other workers
before
+ # releasing the distributed lock. Query execution already commits
its
+ # audit/progress state, so this does not introduce a new
transaction
+ # boundary for the query paths using forced refresh.
+ token_session.commit() # pylint:
disable=consider-using-transaction
return token.access_token
+def execute_with_oauth2_retry(
+ database: Database,
+ operation: Callable[[], T],
+ can_retry: Callable[[], bool] | None = None,
+) -> T:
+ """Refresh a rejected access token and retry an operation once."""
+ # pylint: disable=import-outside-toplevel
+ from superset.models.core import DatabaseUserOAuth2Tokens
+
+ user = getattr(g, "user", None)
+ user_id = getattr(user, "id", None)
+ rejected_access_token = None
+ if user_id is not None:
+ with db.session.no_autoflush:
+ token = (
+ db.session.query(DatabaseUserOAuth2Tokens)
+ .filter_by(user_id=user_id, database_id=database.id)
+ .one_or_none()
+ )
+ rejected_access_token = token.access_token if token is not None else
None
+
+ retry_context = _oauth2_retry_active.set(True)
+ try:
+ try:
+ return operation()
+ finally:
+ _oauth2_retry_active.reset(retry_context)
+ except Exception as ex:
+ is_oauth2_error = (
+ database.is_oauth2_enabled() and
database.db_engine_spec.needs_oauth2(ex)
+ )
+ if not is_oauth2_error:
+ raise
+ if can_retry is not None and not can_retry():
Review Comment:
For a multi-statement SQL Lab query, the first statement commits progress,
so this branch re-raises the OAuth failure after the inner execution path
suppressed `start_oauth2_dance()` for the retry. The query then fails without
the reauthorization URL the user needs. Could this start the OAuth2 dance
before raising (with a regression test for the progressed-query path)?
--
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]
