aminghadersohi commented on code in PR #43623:
URL: https://github.com/apache/superset/pull/43623#discussion_r3900476514
##########
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:
Fixed in ccd0df22ca. Progressed queries now start the OAuth2 dance before
surfacing the non-replayable auth failure, with regression coverage.
--
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]
