diff --git a/airbyte_cdk/sources/streams/call_rate.py b/airbyte_cdk/sources/streams/call_rate.py index 4a06db3b22..ea3e965cab 100644 --- a/airbyte_cdk/sources/streams/call_rate.py +++ b/airbyte_cdk/sources/streams/call_rate.py @@ -555,8 +555,13 @@ def __init__( :param policies: list of policies in this budget :param maximum_attempts_to_acquire: number of attempts before throwing hit ratelimit exception, we put some big number here to avoid situations when many threads compete with each other for a few lots over a significant amount of time + :raises: ValueError if maximum_attempts_to_acquire is less than 1 """ + if maximum_attempts_to_acquire < 1: + raise ValueError( + f"maximum_attempts_to_acquire must be at least 1, got {maximum_attempts_to_acquire}" + ) self._policies = policies self._maximum_attempts_to_acquire = maximum_attempts_to_acquire @@ -628,13 +633,15 @@ def _do_acquire( last_exception = None endpoint = self._extract_endpoint(request) # sometimes we spend all budget before a second attempt, so we have a few more attempts - for attempt in range(1, self._maximum_attempts_to_acquire): + for attempt in range(1, self._maximum_attempts_to_acquire + 1): try: weight = policy.get_weight(request) if isinstance(policy, BaseCallRatePolicy) else 1 policy.try_acquire(request, weight=weight) return except CallRateLimitHit as exc: last_exception = exc + if attempt == self._maximum_attempts_to_acquire: + break if block: if timeout is not None: time_to_wait = min(timedelta(seconds=timeout), exc.time_to_wait) diff --git a/unit_tests/sources/streams/test_call_rate.py b/unit_tests/sources/streams/test_call_rate.py index a423fe5737..265687aa60 100644 --- a/unit_tests/sources/streams/test_call_rate.py +++ b/unit_tests/sources/streams/test_call_rate.py @@ -13,6 +13,7 @@ from airbyte_cdk.models import SyncMode from airbyte_cdk.sources.streams.call_rate import ( + AbstractCallRatePolicy, APIBudget, CallRateLimitHit, FixedWindowCallRatePolicy, @@ -214,6 +215,57 @@ def test_update(self): policy.update(available_calls=10, call_reset_ts=None) +class TestAPIBudgetMaximumAttempts: + @staticmethod + def _exhausted_policy() -> FixedWindowCallRatePolicy: + policy = FixedWindowCallRatePolicy( + matchers=[], + next_reset_ts=datetime.now() + timedelta(hours=1), + period=timedelta(hours=1), + call_limit=1, + ) + policy.try_acquire(Request("GET", "https://example.com"), weight=1) + return policy + + @pytest.mark.parametrize("maximum_attempts_to_acquire", [0, -1]) + def test_rejects_values_below_one(self, maximum_attempts_to_acquire): + with pytest.raises(ValueError, match="maximum_attempts_to_acquire must be at least 1"): + APIBudget(policies=[], maximum_attempts_to_acquire=maximum_attempts_to_acquire) + + def test_single_attempt_non_blocking_raises_when_exhausted(self): + budget = APIBudget(policies=[self._exhausted_policy()], maximum_attempts_to_acquire=1) + with pytest.raises(CallRateLimitHit): + budget.acquire_call(Request("GET", "https://example.com"), block=False) + + def test_single_attempt_blocking_raises_without_sleeping(self, mocker): + sleep_mock = mocker.patch("airbyte_cdk.sources.streams.call_rate.time.sleep") + budget = APIBudget(policies=[self._exhausted_policy()], maximum_attempts_to_acquire=1) + with pytest.raises(CallRateLimitHit): + budget.acquire_call(Request("GET", "https://example.com"), block=True) + sleep_mock.assert_not_called() + + @pytest.mark.parametrize("maximum_attempts_to_acquire", [1, 2, 5]) + def test_try_acquire_called_exactly_maximum_attempts_times( + self, mocker, maximum_attempts_to_acquire + ): + sleep_mock = mocker.patch("airbyte_cdk.sources.streams.call_rate.time.sleep") + policy = mocker.Mock(spec=AbstractCallRatePolicy) + policy.matches.return_value = True + policy.try_acquire.side_effect = CallRateLimitHit( + error="limit hit", item=None, weight=1, rate="1/s", time_to_wait=timedelta(seconds=1) + ) + budget = APIBudget( + policies=[policy], maximum_attempts_to_acquire=maximum_attempts_to_acquire + ) + + with pytest.raises(CallRateLimitHit): + budget.acquire_call(Request("GET", "https://example.com"), block=True) + + assert policy.try_acquire.call_count == maximum_attempts_to_acquire + # no sleep after the final failed attempt + assert sleep_mock.call_count == maximum_attempts_to_acquire - 1 + + class TestFixedWindowCallRatePolicy: def test_limit_rate(self, mocker): policy = FixedWindowCallRatePolicy(