From f10b555703ff61c57b91fa5120711f93186602be Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 13 Oct 2026 07:53:23 +0000 Subject: [PATCH] fix(oauth): surface token refresh retry exhaustion as transient_error Co-Authored-By: bot_apk --- .../concurrent_read_processor.py | 11 ++- .../requests_native_auth/abstract_oauth.py | 54 +++++++---- .../sources/declarative/auth/test_oauth.py | 25 +++++ .../test_concurrent_read_processor.py | 41 ++++++++ .../test_requests_native_auth.py | 95 +++++++++++++++++++ 5 files changed, 207 insertions(+), 19 deletions(-) diff --git a/airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py b/airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py index 160ec1ef03..d5a70760ec 100644 --- a/airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py +++ b/airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py @@ -32,6 +32,7 @@ from airbyte_cdk.sources.streams.concurrent.partition_reader import PartitionReader from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition from airbyte_cdk.sources.streams.concurrent.partitions.types import PartitionCompleteSentinel +from airbyte_cdk.sources.streams.http.exceptions import BaseBackoffException from airbyte_cdk.sources.types import Record from airbyte_cdk.sources.utils.record_helper import stream_data_to_airbyte_message from airbyte_cdk.sources.utils.slice_logger import SliceLogger @@ -252,10 +253,18 @@ def on_exception(self, exception: StreamThreadException) -> Iterable[AirbyteMess if isinstance(exception.exception, AirbyteTracedException): yield exception.exception.as_airbyte_message(stream_descriptor=stream_descriptor) else: + inner_exception = exception.exception + failure_type = ( + inner_exception.failure_type + if isinstance(inner_exception, BaseBackoffException) + and inner_exception.failure_type is not None + else FailureType.system_error + ) yield AirbyteTracedException.from_exception( - exception.exception, + inner_exception, stream_descriptor=stream_descriptor, message=f"An unexpected error occurred in stream {exception.stream_name}: {type(exception.exception).__name__}", + failure_type=failure_type, ).as_airbyte_message() def _flag_exception(self, stream_name: str, exception: Exception) -> None: 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 8fa300b176..ec592c3aa9 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 @@ -409,24 +409,8 @@ def _wrap_refresh_token_exception( ), max_time=300, ) - def _make_handled_request(self) -> Any: - """ - Makes a handled HTTP request to refresh an OAuth token. - - This method sends a POST request to the token refresh endpoint with the necessary - headers and body to obtain a new access token. It handles various exceptions that - may occur during the request and logs the response for troubleshooting purposes. - - Returns: - Mapping[str, Any]: The JSON response from the token refresh endpoint. - - Raises: - DefaultBackoffException: If the response status code is 429 (Too Many Requests) - or any 5xx server error. - AirbyteTracedException: If the refresh token is invalid or expired, prompting - re-authentication. - Exception: For any other exceptions that occur during the request. - """ + def _make_handled_request_with_backoff(self) -> Any: + """Sends the token refresh request, retrying 429/5xx and network errors with backoff.""" try: response = requests.request( method="POST", @@ -492,6 +476,40 @@ def _make_handled_request(self) -> Any: failure_type=FailureType.system_error, ) from e + def _make_handled_request(self) -> Any: + """ + Makes a handled HTTP request to refresh an OAuth token. + + This method sends a POST request to the token refresh endpoint with the necessary + headers and body to obtain a new access token. It handles various exceptions that + may occur during the request and logs the response for troubleshooting purposes. + + Returns: + Mapping[str, Any]: The JSON response from the token refresh endpoint. + + Raises: + AirbyteTracedException: With `transient_error` when 429/5xx retries are + exhausted, or with `config_error` if the refresh + token is invalid or expired, prompting + re-authentication. + Exception: For any other exceptions that occur during the request. + """ + try: + return self._make_handled_request_with_backoff() + except DefaultBackoffException as e: + status_code = e.response.status_code if e.response is not None else None + if status_code == 429: + message = "OAuth token refresh failed because the provider's token endpoint rate limited the request (HTTP 429)." + elif status_code is not None: + message = f"OAuth token refresh failed because the provider's token endpoint was temporarily unavailable (HTTP {status_code})." + else: + message = "OAuth token refresh failed because the provider's token endpoint was temporarily unavailable." + raise AirbyteTracedException( + message=message, + internal_message=f"OAuth token refresh retries were exhausted. {self._build_provider_response_info(e)}", + failure_type=e.failure_type or FailureType.transient_error, + ) from e + def _ensure_access_token_in_response(self, response_data: Mapping[str, Any]) -> None: """ Ensures that the access token is present in the response data. diff --git a/unit_tests/sources/declarative/auth/test_oauth.py b/unit_tests/sources/declarative/auth/test_oauth.py index 71b8328b33..0327859b86 100644 --- a/unit_tests/sources/declarative/auth/test_oauth.py +++ b/unit_tests/sources/declarative/auth/test_oauth.py @@ -17,6 +17,7 @@ from airbyte_cdk.models import FailureType from airbyte_cdk.sources.declarative.auth import DeclarativeOauth2Authenticator from airbyte_cdk.sources.declarative.auth.jwt import JwtAuthenticator +from airbyte_cdk.sources.streams.http.exceptions import DefaultBackoffException from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse from airbyte_cdk.utils import AirbyteTracedException from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets, update_secrets @@ -846,6 +847,30 @@ def test_transient_network_error_wrapped_as_transient_error(self, exception_clas assert exc_info.value.failure_type == FailureType.transient_error assert "network error" in exc_info.value.message.lower() + @pytest.mark.usefixtures("mock_sleep") + def test_refresh_access_token_retries_exhausted_raises_transient_error(self, requests_mock): + """A token endpoint that keeps returning 503 past the backoff window raises AirbyteTracedException with transient_error.""" + oauth = self._create_authenticator() + requests_mock.post( + "https://refresh_endpoint.com", + status_code=503, + text="Server Unavailable", + ) + + with pytest.raises(AirbyteTracedException) as exc_info: + oauth.refresh_access_token() + + exc = exc_info.value + assert exc.failure_type == FailureType.transient_error + assert ( + exc.message + == "OAuth token refresh failed because the provider's token endpoint was temporarily unavailable (HTTP 503)." + ) + assert "DefaultBackoffException" not in exc.message + assert "https://" not in exc.message + assert isinstance(exc.__cause__, DefaultBackoffException) + assert requests_mock.call_count > 1 + def test_connection_error_is_retried_before_raising(self, mocker): """ConnectionError triggers backoff retries in _make_handled_request before propagating.""" oauth = self._create_authenticator() diff --git a/unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py b/unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py index 7b38486abe..bee35242f2 100644 --- a/unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py +++ b/unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py @@ -7,6 +7,7 @@ import freezegun import pytest +import requests from airbyte_cdk.models import ( AirbyteLogMessage, @@ -48,6 +49,7 @@ from airbyte_cdk.sources.streams.concurrent.partition_reader import PartitionReader from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition from airbyte_cdk.sources.streams.concurrent.partitions.types import PartitionCompleteSentinel +from airbyte_cdk.sources.streams.http.exceptions import DefaultBackoffException from airbyte_cdk.sources.types import Record from airbyte_cdk.sources.utils.slice_logger import SliceLogger from airbyte_cdk.utils.traced_exception import AirbyteTracedException @@ -825,6 +827,45 @@ def test_on_exception_non_ate_uses_templated_message_with_correct_failure_type(s # failure_type defaults to system_error for unhandled exceptions assert trace_error.failure_type == FailureType.system_error + def test_on_exception_non_ate_backoff_exception_preserves_failure_type(self): + """Backoff exceptions keep their failure_type when wrapped in on_exception's non-ATE branch.""" + handler = ConcurrentReadProcessor( + [self._stream], + self._partition_enqueuer, + self._thread_pool_manager, + self._logger, + self._slice_logger, + self._message_repository, + self._partition_reader, + ) + + for backoff_failure_type, expected_failure_type in [ + (FailureType.transient_error, FailureType.transient_error), + (FailureType.config_error, FailureType.config_error), + (None, FailureType.system_error), + ]: + with self.subTest(backoff_failure_type=backoff_failure_type): + prepared_request = requests.Request("POST", "https://token.example.com").prepare() + response = requests.Response() + response.status_code = 503 + response.request = prepared_request + inner_exception = DefaultBackoffException( + request=prepared_request, + response=response, + failure_type=backoff_failure_type, + ) + exception = StreamThreadException(inner_exception, _STREAM_NAME) + + exception_messages = list(handler.on_exception(exception)) + assert len(exception_messages) == 1 + + trace_error = exception_messages[0].trace.error + assert ( + trace_error.message + == f"An unexpected error occurred in stream {_STREAM_NAME}: DefaultBackoffException" + ) + assert trace_error.failure_type == expected_failure_type + @freezegun.freeze_time("2020-01-01T00:00:00") def test_start_next_partition_generator(self): stream_instances_to_read_from = [self._stream] 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 923292cc25..6413ed7669 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 @@ -18,6 +18,7 @@ from requests.exceptions import RequestException from airbyte_cdk.models import FailureType, OrchestratorType, Type +from airbyte_cdk.sources.streams.http.exceptions import DefaultBackoffException from airbyte_cdk.sources.streams.http.requests_native_auth import ( BasicHttpAuthenticator, MultipleTokenAuthenticator, @@ -1160,6 +1161,100 @@ def test_refresh_access_token_without_error_fields_keeps_bare_guidance(self, req assert "the grant was revoked" in exc_info.value.internal_message +class TestRefreshAccessTokenRetryExhaustion: + """Token refresh endpoints that keep returning 429/5xx past the backoff + window must surface a `transient_error` AirbyteTracedException. + + Not in the freezegun-frozen test classes above: `mock_sleep` advances the + backoff clock by ticking a frozen datetime, and a class-level freezer + shadows that tick so the 300s window would never elapse. + """ + + @pytest.mark.usefixtures("mock_sleep") + @pytest.mark.parametrize( + "error_code, expected_message", + [ + ( + 503, + "OAuth token refresh failed because the provider's token endpoint was temporarily unavailable (HTTP 503).", + ), + ( + 500, + "OAuth token refresh failed because the provider's token endpoint was temporarily unavailable (HTTP 500).", + ), + ( + 429, + "OAuth token refresh failed because the provider's token endpoint rate limited the request (HTTP 429).", + ), + ], + ) + def test_refresh_access_token_retries_exhausted_raises_transient_error( + self, error_code, expected_message, requests_mock + ): + oauth = Oauth2Authenticator( + f"https://{TestOauth2Authenticator.refresh_endpoint}", + TestOauth2Authenticator.client_id, + TestOauth2Authenticator.client_secret, + TestOauth2Authenticator.refresh_token, + ) + requests_mock.post( + f"https://{TestOauth2Authenticator.refresh_endpoint}", + status_code=error_code, + text="Server Unavailable", + ) + + with pytest.raises(AirbyteTracedException) as exc_info: + oauth.refresh_access_token() + + exc = exc_info.value + assert exc.failure_type == FailureType.transient_error + assert exc.message == expected_message + assert "DefaultBackoffException" not in exc.message + assert "https://" not in exc.message + assert "DOCTYPE" not in exc.message + assert f"HTTP {error_code}" in exc.internal_message + assert "DOCTYPE" in exc.internal_message + assert isinstance(exc.__cause__, DefaultBackoffException) + assert requests_mock.call_count > 1 + + @pytest.mark.usefixtures("mock_sleep") + def test_single_use_refresh_token_retries_exhausted_raises_transient_error(self, requests_mock): + connector_config = { + "credentials": { + "access_token": "my_access_token", + "refresh_token": "my_refresh_token", + "client_id": "my_client_id", + "client_secret": "my_client_secret", + "token_expiry_date": "2022-12-31T00:00:00+00:00", + } + } + authenticator = SingleUseRefreshTokenOauth2Authenticator( + connector_config, + token_refresh_endpoint="https://refresh_endpoint.com", + client_id=connector_config["credentials"]["client_id"], + client_secret=connector_config["credentials"]["client_secret"], + ) + requests_mock.post( + "https://refresh_endpoint.com", + status_code=503, + text="Server Unavailable", + ) + + with pytest.raises(AirbyteTracedException) as exc_info: + authenticator.refresh_access_token() + + exc = exc_info.value + assert exc.failure_type == FailureType.transient_error + assert ( + exc.message + == "OAuth token refresh failed because the provider's token endpoint was temporarily unavailable (HTTP 503)." + ) + assert "DefaultBackoffException" not in exc.message + assert "https://" not in exc.message + assert isinstance(exc.__cause__, DefaultBackoffException) + assert requests_mock.call_count > 1 + + @freezegun.freeze_time("2022-12-31") class TestSingleUseRefreshTokenOauth2Authenticator: @pytest.fixture