diff --git a/airbyte_cdk/sources/streams/http/http_client.py b/airbyte_cdk/sources/streams/http/http_client.py index 8da499aa3..8b1864ff9 100644 --- a/airbyte_cdk/sources/streams/http/http_client.py +++ b/airbyte_cdk/sources/streams/http/http_client.py @@ -377,7 +377,7 @@ def _send_with_retry( if is_rate_limited: raise AirbyteTracedException( internal_message=f"Rate limit retry budget exhausted. Last exception: {e}", - message="API rate limit exceeded.", + message=f"API rate limit exceeded for stream '{self._name}' after all retries were exhausted.", failure_type=FailureType.transient_error, exception=e, stream_descriptor=StreamDescriptor(name=self._name), diff --git a/unit_tests/sources/streams/http/test_http.py b/unit_tests/sources/streams/http/test_http.py index 7f4f9ac41..356616a48 100644 --- a/unit_tests/sources/streams/http/test_http.py +++ b/unit_tests/sources/streams/http/test_http.py @@ -4,6 +4,7 @@ import json import logging +import re from http import HTTPStatus from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union from unittest.mock import ANY, MagicMock, Mock, patch @@ -223,7 +224,9 @@ def get_error_handler(self) -> Optional[ErrorHandler]: with pytest.raises( AirbyteTracedException, - match="API rate limit exceeded.", + match=re.escape( + f"API rate limit exceeded for stream '{stream.name}' after all retries were exhausted." + ), ): list(stream.read_records(SyncMode.full_refresh)) if retries <= 0: @@ -317,12 +320,33 @@ def test_raise_on_http_errors_off_429(mocker): mocker.patch.object(requests.Session, "send", return_value=req) with pytest.raises( AirbyteTracedException, - match="API rate limit exceeded.", + match=re.escape( + f"API rate limit exceeded for stream '{stream.name}' after all retries were exhausted." + ), ): stream.exit_on_rate_limit = True list(stream.read_records(SyncMode.full_refresh)) +def test_rate_limit_exhaustion_includes_stream_context(mocker): + mocker.patch("time.sleep", lambda x: None) + stream = AutoFailFalseHttpStream() + req = requests.Response() + req.status_code = 429 + mocker.patch.object(requests.Session, "send", return_value=req) + + stream.exit_on_rate_limit = True + with pytest.raises(AirbyteTracedException) as exc_info: + list(stream.read_records(SyncMode.full_refresh)) + + assert ( + exc_info.value.message + == f"API rate limit exceeded for stream '{stream.name}' after all retries were exhausted." + ) + assert exc_info.value.failure_type == FailureType.transient_error + assert exc_info.value._stream_descriptor.name == stream.name + + @pytest.mark.parametrize("status_code", [500, 501, 503, 504]) def test_raise_on_http_errors_off_5xx(mocker, status_code): mocker.patch("time.sleep", lambda x: None) diff --git a/unit_tests/sources/streams/http/test_http_client.py b/unit_tests/sources/streams/http/test_http_client.py index 31d4ff9a1..bdb821861 100644 --- a/unit_tests/sources/streams/http/test_http_client.py +++ b/unit_tests/sources/streams/http/test_http_client.py @@ -643,6 +643,30 @@ def backoff_time(self, *args, **kwargs): assert mocked_send.call_count == retries + 1 +@pytest.mark.usefixtures("mock_sleep") +def test_rate_limit_exhaustion_message_includes_stream_name(): + stream_name = "job extract_url - my_stream" + mocked_session = MagicMock(spec=requests.Session) + mocked_response = requests.Response() + mocked_response.status_code = 429 + mocked_session.send.return_value = mocked_response + http_client = HttpClient( + name=stream_name, + logger=MagicMock(), + error_handler=HttpStatusErrorHandler(logger=MagicMock(), max_retries=1), + session=mocked_session, + ) + prepared_request = requests.Request("GET", "https://test_base_url.com/endpoint").prepare() + + with pytest.raises(AirbyteTracedException) as exc_info: + http_client._send_with_retry(prepared_request, request_kwargs={}) + + assert ( + exc_info.value.message + == f"API rate limit exceeded for stream '{stream_name}' after all retries were exhausted." + ) + + @pytest.mark.usefixtures("mock_sleep") def test_backoff_strategy_max_time(): error_handler = HttpStatusErrorHandler(