From cc554e244a3ccb9e1200cc88f543c8a1d41ba14f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 27 Jan 2026 12:02:20 +0000 Subject: [PATCH 1/2] fix(oauth): add thread-safe token refresh with double-checked locking When multiple streams start concurrently with an expired token, they can all detect the token is expired and attempt to refresh simultaneously. For single-use refresh tokens, this causes failures because the refresh token is invalidated after the first successful refresh. This fix adds a class-level lock to AbstractOauth2Authenticator that serializes token refresh attempts. The double-checked locking pattern ensures that: 1. Only one thread performs the actual refresh 2. Other threads wait and use the newly refreshed token 3. No redundant refresh attempts occur Changes: - Add _token_refresh_lock class-level threading.Lock to AbstractOauth2Authenticator - Update get_access_token() in AbstractOauth2Authenticator to use double-checked locking - Update get_access_token() in SingleUseRefreshTokenOauth2Authenticator with same pattern - Add unit tests for concurrent token refresh scenarios Co-Authored-By: gl_anatolii.yatsuk@airbyte.io --- .../requests_native_auth/abstract_oauth.py | 26 +++- .../http/requests_native_auth/oauth.py | 29 ++-- .../test_requests_native_auth.py | 135 ++++++++++++++++++ 3 files changed, 177 insertions(+), 13 deletions(-) diff --git a/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py b/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py index 3b4aa98442..9f69552ece 100644 --- a/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py +++ b/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py @@ -3,6 +3,7 @@ # import logging +import threading from abc import abstractmethod from datetime import timedelta from json import JSONDecodeError @@ -41,6 +42,13 @@ class AbstractOauth2Authenticator(AuthBase): _NO_STREAM_NAME = None + # Class-level lock to prevent concurrent token refresh across multiple authenticator instances. + # This is necessary because multiple streams may share the same OAuth credentials (refresh token) + # through the connector config. Without this lock, concurrent refresh attempts can cause race + # conditions where one stream successfully refreshes the token while others fail because the + # refresh token has been invalidated (especially for single-use refresh tokens). + _token_refresh_lock: threading.Lock = threading.Lock() + def __init__( self, refresh_token_error_status_codes: Tuple[int, ...] = (), @@ -86,11 +94,21 @@ def get_auth_header(self) -> Mapping[str, Any]: return {"Authorization": f"Bearer {token}"} def get_access_token(self) -> str: - """Returns the access token""" + """ + Returns the access token. + + This method uses double-checked locking to ensure thread-safe token refresh. + When multiple threads (streams) detect an expired token simultaneously, only one + will perform the refresh while others wait. After acquiring the lock, the token + expiry is re-checked to avoid redundant refresh attempts. + """ if self.token_has_expired(): - token, expires_in = self.refresh_access_token() - self.access_token = token - self.set_token_expiry_date(expires_in) + with self._token_refresh_lock: + # Double-check after acquiring lock - another thread may have already refreshed + if self.token_has_expired(): + token, expires_in = self.refresh_access_token() + self.access_token = token + self.set_token_expiry_date(expires_in) return self.access_token diff --git a/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py b/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py index a2932c2949..479ba06a32 100644 --- a/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py +++ b/airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py @@ -317,19 +317,30 @@ def token_has_expired(self) -> bool: return ab_datetime_now() > self.get_token_expiry_date() def get_access_token(self) -> str: - """Retrieve new access and refresh token if the access token has expired. - The new refresh token is persisted with the set_refresh_token function + """ + Retrieve new access and refresh token if the access token has expired. + + This method uses double-checked locking to ensure thread-safe token refresh. + This is especially critical for single-use refresh tokens where concurrent + refresh attempts would cause failures as the refresh token is invalidated + after first use. + + The new refresh token is persisted with the set_refresh_token function. + Returns: str: The current access_token, updated if it was previously expired. """ if self.token_has_expired(): - new_access_token, access_token_expires_in, new_refresh_token = ( - self.refresh_access_token() - ) - self.access_token = new_access_token - self.set_refresh_token(new_refresh_token) - self.set_token_expiry_date(access_token_expires_in) - self._emit_control_message() + with self._token_refresh_lock: + # Double-check after acquiring lock - another thread may have already refreshed + if self.token_has_expired(): + new_access_token, access_token_expires_in, new_refresh_token = ( + self.refresh_access_token() + ) + self.access_token = new_access_token + self.set_refresh_token(new_refresh_token) + self.set_token_expiry_date(access_token_expires_in) + self._emit_control_message() return self.access_token def refresh_access_token(self) -> Tuple[str, AirbyteDateTime, str]: # type: ignore[override] diff --git a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py index dbfc0ac866..4c71465c3d 100644 --- a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py +++ b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py @@ -785,3 +785,138 @@ def mock_request(method, url, data, headers): raise Exception( f"Error while refreshing access token with request: {method}, {url}, {data}, {headers}" ) + + +class TestConcurrentTokenRefresh: + """ + Test class for verifying thread-safe token refresh behavior. + + These tests ensure that when multiple threads (streams) attempt to refresh + an expired token simultaneously, only one refresh actually occurs and + others wait and use the refreshed token. + """ + + def test_concurrent_token_refresh_only_refreshes_once(self, mocker): + """ + When multiple threads detect an expired token and try to refresh simultaneously, + only one thread should actually perform the refresh. Others should wait and + use the newly refreshed token. + """ + import threading + import time + + refresh_call_count = 0 + refresh_call_lock = threading.Lock() + + def mock_refresh_access_token(self): + nonlocal refresh_call_count + with refresh_call_lock: + refresh_call_count += 1 + time.sleep(0.1) + return ("new_access_token", ab_datetime_now() + timedelta(hours=1)) + + mocker.patch.object( + Oauth2Authenticator, + "refresh_access_token", + mock_refresh_access_token, + ) + + oauth = Oauth2Authenticator( + token_refresh_endpoint="https://refresh_endpoint.com", + client_id="client_id", + client_secret="client_secret", + refresh_token="refresh_token", + token_expiry_date=ab_datetime_now() - timedelta(hours=1), + ) + + results = [] + errors = [] + + def get_token(): + try: + token = oauth.get_access_token() + results.append(token) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=get_token) for _ in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(errors) == 0, f"Unexpected errors: {errors}" + assert len(results) == 5 + assert all(token == "new_access_token" for token in results) + assert refresh_call_count == 1, f"Expected 1 refresh call, got {refresh_call_count}" + + def test_single_use_refresh_token_concurrent_refresh_only_refreshes_once(self, mocker): + """ + For SingleUseRefreshTokenOauth2Authenticator, concurrent refresh attempts + should also only result in one actual refresh to prevent invalidating + the single-use refresh token. + """ + import threading + import time + + refresh_call_count = 0 + refresh_call_lock = threading.Lock() + + connector_config = { + "credentials": { + "client_id": "client_id", + "client_secret": "client_secret", + "refresh_token": "refresh_token", + "access_token": "old_access_token", + "token_expiry_date": str(ab_datetime_now() - timedelta(hours=1)), + } + } + + def mock_refresh_access_token(self): + nonlocal refresh_call_count + with refresh_call_lock: + refresh_call_count += 1 + time.sleep(0.1) + return ( + "new_access_token", + ab_datetime_now() + timedelta(hours=1), + "new_refresh_token", + ) + + mocker.patch.object( + SingleUseRefreshTokenOauth2Authenticator, + "refresh_access_token", + mock_refresh_access_token, + ) + + mocker.patch.object( + SingleUseRefreshTokenOauth2Authenticator, + "_emit_control_message", + lambda self: None, + ) + + oauth = SingleUseRefreshTokenOauth2Authenticator( + connector_config=connector_config, + token_refresh_endpoint="https://refresh_endpoint.com", + ) + + results = [] + errors = [] + + def get_token(): + try: + token = oauth.get_access_token() + results.append(token) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target=get_token) for _ in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(errors) == 0, f"Unexpected errors: {errors}" + assert len(results) == 5 + assert all(token == "new_access_token" for token in results) + assert refresh_call_count == 1, f"Expected 1 refresh call, got {refresh_call_count}" From c88189396ab90733f8ada1eb3e6481b47efb8ecc Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 27 Jan 2026 14:25:59 +0000 Subject: [PATCH 2/2] refactor: move threading and time imports to top of test file Address PR review feedback from aldogonzalez8 to follow PEP8 import conventions. Co-Authored-By: gl_anatolii.yatsuk@airbyte.io --- .../requests_native_auth/test_requests_native_auth.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py index 4c71465c3d..71183b5aad 100644 --- a/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py +++ b/unit_tests/sources/streams/http/requests_native_auth/test_requests_native_auth.py @@ -4,6 +4,8 @@ import json import logging +import threading +import time from datetime import timedelta from typing import Optional, Union from unittest.mock import Mock @@ -802,9 +804,6 @@ def test_concurrent_token_refresh_only_refreshes_once(self, mocker): only one thread should actually perform the refresh. Others should wait and use the newly refreshed token. """ - import threading - import time - refresh_call_count = 0 refresh_call_lock = threading.Lock() @@ -856,9 +855,6 @@ def test_single_use_refresh_token_concurrent_refresh_only_refreshes_once(self, m should also only result in one actual refresh to prevent invalidating the single-use refresh token. """ - import threading - import time - refresh_call_count = 0 refresh_call_lock = threading.Lock()