From 4759654e5f840d3eaf82c9e71bd1eb6bfe361b1a Mon Sep 17 00:00:00 2001 From: Artem Inzhyyants <36314070+artem1205@users.noreply.github.com> Date: Wed, 18 Dec 2024 16:48:57 +0100 Subject: [PATCH] fix: (HttpClient) rate limit fix unlimited tries (#171) Signed-off-by: Artem Inzhyyants --- airbyte_cdk/sources/streams/http/http_client.py | 6 ++++-- unit_tests/sources/streams/http/test_http_client.py | 7 +++++-- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/airbyte_cdk/sources/streams/http/http_client.py b/airbyte_cdk/sources/streams/http/http_client.py index 715484410..c4fa86866 100644 --- a/airbyte_cdk/sources/streams/http/http_client.py +++ b/airbyte_cdk/sources/streams/http/http_client.py @@ -262,7 +262,7 @@ def _send_with_retry( user_backoff_handler = user_defined_backoff_handler(max_tries=max_tries, max_time=max_time)( self._send ) - rate_limit_backoff_handler = rate_limit_default_backoff_handler() + rate_limit_backoff_handler = rate_limit_default_backoff_handler(max_tries=max_tries) backoff_handler = http_client_default_backoff_handler( max_tries=max_tries, max_time=max_time ) @@ -472,7 +472,9 @@ def _handle_error_resolution( elif retry_endlessly: raise RateLimitBackoffException( - request=request, response=response or exc, error_message=error_message + request=request, + response=(response if response is not None else exc), + error_message=error_message, ) raise DefaultBackoffException( diff --git a/unit_tests/sources/streams/http/test_http_client.py b/unit_tests/sources/streams/http/test_http_client.py index 29bac0ec8..5cc6d20e4 100644 --- a/unit_tests/sources/streams/http/test_http_client.py +++ b/unit_tests/sources/streams/http/test_http_client.py @@ -20,6 +20,7 @@ ) from airbyte_cdk.sources.streams.http.exceptions import ( DefaultBackoffException, + RateLimitBackoffException, RequestBodyException, UserDefinedBackoffException, ) @@ -690,10 +691,12 @@ def backoff_time(self, *args, **kwargs): @pytest.mark.parametrize( "exit_on_rate_limit, expected_call_count, expected_error", - [[True, 6, DefaultBackoffException], [False, 38, OverflowError]], + [[True, 6, DefaultBackoffException], [False, 6, RateLimitBackoffException]], ) @pytest.mark.usefixtures("mock_sleep") -def test_backoff_strategy_endless(exit_on_rate_limit, expected_call_count, expected_error): +def test_backoff_strategy_endless( + exit_on_rate_limit: bool, expected_call_count: int, expected_error: Exception +): http_client = HttpClient( name="test", logger=MagicMock(), error_handler=HttpStatusErrorHandler(logger=MagicMock()) )