aminghadersohi commented on code in PR #44765:
URL: https://github.com/apache/superset/pull/44765#discussion_r4140678229
##########
superset/utils/oauth2.py:
##########
@@ -322,12 +423,11 @@ def execute_with_oauth2_retry( # noqa: C901
database.db_engine_spec.engine,
)
try:
- access_token = refresh_oauth2_token(
+ access_token = force_refresh_oauth2_token(
Review Comment:
Fixed in 54572820a6. `execute_with_oauth2_retry` now treats
`OAuth2TokenRefreshError` as terminal: it increments
`oauth2.forced_refresh.exchange_rejected`, starts the sign-in dance and
re-raises, so it never opens the isolated session or presents the refused
refresh token a second time.
`test_connection_does_not_re_exchange_a_refused_refresh_token` (parametrized
over `get_raw_connection` and `get_inspector`) starts from an expired stored
token with the provider refusing the refresh. It asserts one exchange, the
deletion flushed but not committed, no isolated `Session`, one
`start_oauth2_dance` and an `OAuth2RedirectError`. Both cases fail with the
guard removed.
##########
superset/utils/oauth2.py:
##########
@@ -268,6 +268,107 @@ def _refresh_oauth2_token_locked( # noqa: C901
return token.access_token
[email protected]_exception(
+ backoff.expo,
+ AcquireDistributedLockFailedException,
+ factor=0.1,
+ base=2,
+ max_tries=8,
Review Comment:
Fixed in 54572820a6. Both decorators read `OAUTH2_LOCK_BACKOFF_FACTOR`,
`OAUTH2_LOCK_BACKOFF_BASE` and `OAUTH2_LOCK_BACKOFF_MAX_TRIES`, so the
expired-token and forced-refresh paths can't drift apart.
##########
tests/unit_tests/utils/oauth2_tests.py:
##########
@@ -907,3 +910,127 @@ def test_get_oauth2_redirect_uri_raises_on_runtime_error(
)
with pytest.raises(OAuth2Error):
get_oauth2_redirect_uri()
+
+
+def test_force_refresh_waits_out_lock_contention(
+ mocker: MockerFixture,
+) -> None:
+ """
+ A loser of the refresh race reuses the winner's token instead of raising.
+
+ Every chart on a dashboard opening a connection with the same rejected
token
+ races for the non-blocking refresh lock. The losers must retry until the
winner
+ commits, then short circuit on the committed access token.
+ """
+ mocker.patch("time.sleep") # avoid backoff delays in tests
+ db = mocker.patch("superset.utils.oauth2.db")
+ mocker.patch("superset.utils.oauth2.Session", return_value=db.session)
+ lock = mocker.patch("superset.utils.oauth2.DistributedLock")
+ lock.side_effect = [
+ LockAlreadyHeldException("Lock already taken"),
+ mocker.MagicMock(),
+ ]
+ db_engine_spec = mocker.MagicMock()
+ token = mocker.MagicMock(access_token="winning-token") # noqa: S106
+
db.session.query().populate_existing().filter_by().one_or_none.return_value =
token
+
+ result = force_refresh_oauth2_token(
+ DUMMY_OAUTH2_CONFIG,
+ 1,
+ 2,
+ db_engine_spec,
+ rejected_access_token="rejected-token", # noqa: S106
+ )
+
+ assert result == "winning-token"
+ assert lock.call_count == 2
+ db_engine_spec.get_oauth2_fresh_token.assert_not_called()
+
+
+def test_force_refresh_reads_committed_token_when_lock_never_frees(
+ mocker: MockerFixture,
+) -> None:
+ """Giving up on the lock still reuses a token another worker committed."""
+ mocker.patch("time.sleep") # avoid backoff delays in tests
+ db = mocker.patch("superset.utils.oauth2.db")
+ mocker.patch("superset.utils.oauth2.Session", return_value=db.session)
+ mocker.patch(
+ "superset.utils.oauth2.DistributedLock",
+ side_effect=AcquireDistributedLockFailedException("Lock not
available"),
+ )
+ db_engine_spec = mocker.MagicMock()
+ db.session.query().filter_by().one_or_none.return_value = mocker.MagicMock(
+ access_token="winning-token" # noqa: S106
+ )
+
+ result = force_refresh_oauth2_token(
+ DUMMY_OAUTH2_CONFIG,
+ 1,
+ 2,
+ db_engine_spec,
+ rejected_access_token="rejected-token", # noqa: S106
+ )
+
+ assert result == "winning-token"
Review Comment:
This finding's premise is wrong. `raise_on_giveup=False` is set only on
`get_oauth2_access_token`. The forced-refresh helper (renamed
`_retry_forced_refresh_until_lock_available` in 54572820a6, following your
naming suggestion) omits it, so the lock exception reaches
`force_refresh_oauth2_token`'s `except` and `_read_refreshed_access_token`
runs. To pin that, 54572820a6 also asserts that the lock is attempted
`OAUTH2_LOCK_BACKOFF_MAX_TRIES` times and that the fallback's token read
happens.
##########
tests/unit_tests/utils/oauth2_tests.py:
##########
@@ -907,3 +910,127 @@ def test_get_oauth2_redirect_uri_raises_on_runtime_error(
)
with pytest.raises(OAuth2Error):
get_oauth2_redirect_uri()
+
+
+def test_force_refresh_waits_out_lock_contention(
+ mocker: MockerFixture,
+) -> None:
+ """
+ A loser of the refresh race reuses the winner's token instead of raising.
+
+ Every chart on a dashboard opening a connection with the same rejected
token
+ races for the non-blocking refresh lock. The losers must retry until the
winner
+ commits, then short circuit on the committed access token.
+ """
+ mocker.patch("time.sleep") # avoid backoff delays in tests
+ db = mocker.patch("superset.utils.oauth2.db")
+ mocker.patch("superset.utils.oauth2.Session", return_value=db.session)
+ lock = mocker.patch("superset.utils.oauth2.DistributedLock")
+ lock.side_effect = [
+ LockAlreadyHeldException("Lock already taken"),
+ mocker.MagicMock(),
+ ]
+ db_engine_spec = mocker.MagicMock()
+ token = mocker.MagicMock(access_token="winning-token") # noqa: S106
+
db.session.query().populate_existing().filter_by().one_or_none.return_value =
token
+
+ result = force_refresh_oauth2_token(
+ DUMMY_OAUTH2_CONFIG,
+ 1,
+ 2,
+ db_engine_spec,
+ rejected_access_token="rejected-token", # noqa: S106
+ )
+
+ assert result == "winning-token"
+ assert lock.call_count == 2
+ db_engine_spec.get_oauth2_fresh_token.assert_not_called()
+
+
+def test_force_refresh_reads_committed_token_when_lock_never_frees(
+ mocker: MockerFixture,
+) -> None:
+ """Giving up on the lock still reuses a token another worker committed."""
+ mocker.patch("time.sleep") # avoid backoff delays in tests
+ db = mocker.patch("superset.utils.oauth2.db")
+ mocker.patch("superset.utils.oauth2.Session", return_value=db.session)
+ mocker.patch(
+ "superset.utils.oauth2.DistributedLock",
+ side_effect=AcquireDistributedLockFailedException("Lock not
available"),
+ )
+ db_engine_spec = mocker.MagicMock()
+ db.session.query().filter_by().one_or_none.return_value = mocker.MagicMock(
+ access_token="winning-token" # noqa: S106
+ )
+
+ result = force_refresh_oauth2_token(
+ DUMMY_OAUTH2_CONFIG,
+ 1,
+ 2,
+ db_engine_spec,
+ rejected_access_token="rejected-token", # noqa: S106
+ )
+
+ assert result == "winning-token"
+ db_engine_spec.get_oauth2_fresh_token.assert_not_called()
+
+
+def test_force_refresh_returns_none_when_no_one_refreshed(
+ mocker: MockerFixture,
+) -> None:
+ """With the lock held and the stored token unchanged there is nothing to
reuse."""
+ mocker.patch("time.sleep") # avoid backoff delays in tests
+ db = mocker.patch("superset.utils.oauth2.db")
+ mocker.patch("superset.utils.oauth2.Session", return_value=db.session)
+ mocker.patch(
+ "superset.utils.oauth2.DistributedLock",
+ side_effect=LockAlreadyHeldException("Lock already taken"),
+ )
+ db.session.query().filter_by().one_or_none.return_value = mocker.MagicMock(
+ access_token="rejected-token" # noqa: S106
+ )
+
+ result = force_refresh_oauth2_token(
+ DUMMY_OAUTH2_CONFIG,
+ 1,
+ 2,
+ mocker.MagicMock(),
+ rejected_access_token="rejected-token", # noqa: S106
+ )
+
+ assert result is None
Review Comment:
This finding's premise is wrong. `raise_on_giveup=False` is set only on
`get_oauth2_access_token`. The forced-refresh helper (renamed
`_retry_forced_refresh_until_lock_available` in 54572820a6, following your
naming suggestion) omits it, so the lock exception reaches
`force_refresh_oauth2_token`'s `except` and `_read_refreshed_access_token`
runs. To pin that, 54572820a6 also asserts that the lock is attempted
`OAUTH2_LOCK_BACKOFF_MAX_TRIES` times and that the fallback's token read
happens.
##########
tests/unit_tests/utils/oauth2_tests.py:
##########
@@ -907,3 +910,127 @@ def test_get_oauth2_redirect_uri_raises_on_runtime_error(
)
with pytest.raises(OAuth2Error):
get_oauth2_redirect_uri()
+
+
+def test_force_refresh_waits_out_lock_contention(
+ mocker: MockerFixture,
+) -> None:
+ """
+ A loser of the refresh race reuses the winner's token instead of raising.
+
+ Every chart on a dashboard opening a connection with the same rejected
token
+ races for the non-blocking refresh lock. The losers must retry until the
winner
+ commits, then short circuit on the committed access token.
+ """
+ mocker.patch("time.sleep") # avoid backoff delays in tests
+ db = mocker.patch("superset.utils.oauth2.db")
+ mocker.patch("superset.utils.oauth2.Session", return_value=db.session)
+ lock = mocker.patch("superset.utils.oauth2.DistributedLock")
+ lock.side_effect = [
+ LockAlreadyHeldException("Lock already taken"),
+ mocker.MagicMock(),
+ ]
+ db_engine_spec = mocker.MagicMock()
+ token = mocker.MagicMock(access_token="winning-token") # noqa: S106
+
db.session.query().populate_existing().filter_by().one_or_none.return_value =
token
+
+ result = force_refresh_oauth2_token(
+ DUMMY_OAUTH2_CONFIG,
+ 1,
+ 2,
+ db_engine_spec,
+ rejected_access_token="rejected-token", # noqa: S106
+ )
+
+ assert result == "winning-token"
+ assert lock.call_count == 2
+ db_engine_spec.get_oauth2_fresh_token.assert_not_called()
+
+
+def test_force_refresh_reads_committed_token_when_lock_never_frees(
+ mocker: MockerFixture,
+) -> None:
+ """Giving up on the lock still reuses a token another worker committed."""
+ mocker.patch("time.sleep") # avoid backoff delays in tests
+ db = mocker.patch("superset.utils.oauth2.db")
+ mocker.patch("superset.utils.oauth2.Session", return_value=db.session)
+ mocker.patch(
+ "superset.utils.oauth2.DistributedLock",
+ side_effect=AcquireDistributedLockFailedException("Lock not
available"),
+ )
+ db_engine_spec = mocker.MagicMock()
+ db.session.query().filter_by().one_or_none.return_value = mocker.MagicMock(
+ access_token="winning-token" # noqa: S106
+ )
+
+ result = force_refresh_oauth2_token(
+ DUMMY_OAUTH2_CONFIG,
+ 1,
+ 2,
+ db_engine_spec,
+ rejected_access_token="rejected-token", # noqa: S106
+ )
+
+ assert result == "winning-token"
+ db_engine_spec.get_oauth2_fresh_token.assert_not_called()
+
+
+def test_force_refresh_returns_none_when_no_one_refreshed(
+ mocker: MockerFixture,
+) -> None:
+ """With the lock held and the stored token unchanged there is nothing to
reuse."""
+ mocker.patch("time.sleep") # avoid backoff delays in tests
+ db = mocker.patch("superset.utils.oauth2.db")
+ mocker.patch("superset.utils.oauth2.Session", return_value=db.session)
+ mocker.patch(
+ "superset.utils.oauth2.DistributedLock",
+ side_effect=LockAlreadyHeldException("Lock already taken"),
+ )
+ db.session.query().filter_by().one_or_none.return_value = mocker.MagicMock(
+ access_token="rejected-token" # noqa: S106
+ )
+
+ result = force_refresh_oauth2_token(
+ DUMMY_OAUTH2_CONFIG,
+ 1,
+ 2,
+ mocker.MagicMock(),
+ rejected_access_token="rejected-token", # noqa: S106
+ )
+
+ assert result is None
+
+
+def test_execute_with_oauth2_retry_survives_lock_contention(
+ mocker: MockerFixture,
+) -> None:
+ """
+ A chart that loses the refresh race renders with the winner's token.
+
+ Without contention tolerance the lock failure is not an OAuth2 error, so it
+ escapes as an opaque `AcquireDistributedLockFailedException` and the chart
fails.
+ """
+ mocker.patch("time.sleep") # avoid backoff delays in tests
+ operation = mocker.Mock(side_effect=[RuntimeError("stale OAuth token"),
"result"])
+ database = mocker.MagicMock()
+ database.id = 1
+ database.is_oauth2_enabled.return_value = True
+ database.db_engine_spec.needs_oauth2.return_value = True
+ database.get_oauth2_config.return_value = DUMMY_OAUTH2_CONFIG
+ mocker.patch("superset.utils.oauth2.g").user.id = 2
+ db = mocker.patch("superset.utils.oauth2.db")
+ mocker.patch("superset.utils.oauth2.Session", return_value=db.session)
+ mocker.patch(
+ "superset.utils.oauth2.DistributedLock",
+ side_effect=LockAlreadyHeldException("Lock already taken"),
+ )
+ # The rejected token is read first, then the winner's committed
replacement.
+ db.session.query().filter_by().one_or_none.side_effect = [
+ mocker.MagicMock(access_token="rejected-token"), # noqa: S106
+ mocker.MagicMock(access_token="winning-token"), # noqa: S106
+ ]
+
+ assert execute_with_oauth2_retry(database, operation) == "result"
+
+ assert operation.call_count == 2
+ database.start_oauth2_dance.assert_not_called()
Review Comment:
This finding's premise is wrong. `raise_on_giveup=False` is set only on
`get_oauth2_access_token`. The forced-refresh helper (renamed
`_retry_forced_refresh_until_lock_available` in 54572820a6, following your
naming suggestion) omits it, so the lock exception reaches
`force_refresh_oauth2_token`'s `except` and `_read_refreshed_access_token`
runs. To pin that, 54572820a6 also asserts that the lock is attempted
`OAUTH2_LOCK_BACKOFF_MAX_TRIES` times and that the fallback's token read
happens.
##########
superset/utils/oauth2.py:
##########
@@ -268,6 +268,107 @@ def _refresh_oauth2_token_locked( # noqa: C901
return token.access_token
[email protected]_exception(
+ backoff.expo,
+ AcquireDistributedLockFailedException,
+ factor=0.1,
+ base=2,
+ max_tries=8,
Review Comment:
The mismatch is deliberate, so I've left the window unchanged. The 30s
`ttl_seconds` is a safety ceiling: it frees the lock if its holder dies
mid-exchange. It isn't how long we expect an exchange to take. Waiting out the
TTL would hold a web worker for about 30s on every losing chart request, and
`backoff.expo` uses full jitter by default, so 12.7s is already the upper
bound, not the typical wait. The expired-token path (`get_oauth2_access_token`)
has used the same 12.7s window against the same 30s lock since before this PR.
54572820a6 moves both paths onto shared constants (`OAUTH2_LOCK_BACKOFF_*`), so
they can't diverge if we retune later. A provider that takes 20s for a token
exchange is pathological, and the loser falls back to the sign-in dance as it
did before, so the bounded wait seems the better trade.
--
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]
