aminghadersohi opened a new pull request, #44765: URL: https://github.com/apache/superset/pull/44765
### SUMMARY Some databases (Snowflake, for example) reject an OAuth2 access token when the connection **logs in**. If the token store still holds that token as unexpired, the stored refresh token is never used: the user is sent back to the sign-in page even though a refresh would have worked. `execute_with_oauth2_retry()` (#43623) handles exactly this, but it is only called by the new SQL executor and its Celery task. The paths that open connections elsewhere never use it: - the SQL Lab API (`superset.sql_lab` → `Database.get_raw_connection()`), - chart data (`Database.get_df()` → `get_raw_connection()`, and before that the default schema read for the chart's cache key through `get_inspector()`), - metadata calls through `get_inspector()`. This change routes opening the connection (`get_raw_connection()`) and the inspector (`get_inspector()`) through `execute_with_oauth2_retry()`: - a rejected token triggers one forced refresh, then the connection is opened again; - only the login is retried; work the caller does with the connection is never replayed; - the wrapper steps aside when there is no signed-in user (the existing sign-in handling applies unchanged) or when an outer operation already owns the recovery (the new executor), so a token is exchanged at most once. Engines that reject the token later, per statement (e.g. Google Sheets), are not covered by this change. ### TESTING INSTRUCTIONS New unit tests in `tests/unit_tests/models/core_test.py`: - `test_get_raw_connection_refreshes_a_token_rejected_at_login` - `test_get_inspector_refreshes_a_token_rejected_at_login` - `test_get_raw_connection_asks_to_sign_in_when_the_new_token_is_rejected` - `test_get_raw_connection_does_not_replay_the_callers_block` - `test_get_raw_connection_leaves_recovery_to_an_outer_retry` - `test_get_raw_connection_without_oauth2_is_unchanged` The first three fail without the change; the others guard behaviour that must not change. The existing `tests/unit_tests/models/`, `utils/oauth2_tests.py`, `sql_lab_test.py` and `db_engine_specs/` suites pass. The same approach was verified against a live Snowflake account: the rejected token (error 390303) was exchanged once with the stored refresh token, the new token was persisted, and the connection was opened again with it, for both SQL Lab and chart data; without the change no exchange happened and the user was asked to sign in. ### ADDITIONAL INFORMATION - [ ] Has associated issue: - [ ] Required feature flags: - [ ] Changes UI - [ ] Includes DB Migration - [ ] Introduces new feature or API - [ ] Removes existing feature or API -- 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]
