diff --git a/airbyte_cdk/sources/streams/call_rate.py b/airbyte_cdk/sources/streams/call_rate.py index 4a06db3b22..efde07b61a 100644 --- a/airbyte_cdk/sources/streams/call_rate.py +++ b/airbyte_cdk/sources/streams/call_rate.py @@ -487,14 +487,29 @@ def update( ): # we do our best to sync buckets with API if available_calls == 0: with self._limiter.lock: - items_to_add = self._bucket.count() < self._bucket.rates[0].limit + now: int = TimeClock().now() # type: ignore[no-untyped-call] + items_to_add = self._get_available_calls(now) if items_to_add > 0: - now: int = TimeClock().now() # type: ignore[no-untyped-call] - self._bucket.put(RateItem(name="dummy", timestamp=now, weight=items_to_add)) + added = self._bucket.put( + RateItem(name="dummy", timestamp=now, weight=items_to_add) + ) + if not added: + logger.warning( + "Failed to sync rate limit bucket with API: could not add %s items", + items_to_add, + ) # TODO: add support if needed, it might be that it is not possible to make a good solution for this case # if available_calls is not None and call_reset_ts is not None: # ts = call_reset_ts.timestamp() + def _get_available_calls(self, now: int) -> int: + """Return the number of calls still allowed by the strictest configured rate at the given timestamp (ms).""" + return min( + rate.limit + - sum(1 for item in self._bucket.items if item.timestamp >= now - rate.interval) + for rate in self._bucket.rates + ) + def __str__(self) -> str: """Return a human-friendly description of the moving window rate policy for logging purposes.""" rates_info = ", ".join( diff --git a/unit_tests/sources/streams/test_call_rate.py b/unit_tests/sources/streams/test_call_rate.py index a423fe5737..b14eaf7c48 100644 --- a/unit_tests/sources/streams/test_call_rate.py +++ b/unit_tests/sources/streams/test_call_rate.py @@ -9,6 +9,7 @@ import pytest import requests +from freezegun import freeze_time from requests import Request from airbyte_cdk.models import SyncMode @@ -305,6 +306,82 @@ def test_multiple_limit_rates(self): assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(3600, 0.1) assert str(excinfo.value) == "Bucket for item=call with Rate limit=2/1.0h is already full" + def test_update_available_calls_zero_exhausts_bucket(self): + with freeze_time("2024-01-01 00:00:00") as frozen: + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + for _ in range(3): + policy.try_acquire("call", weight=1) + + policy.update(available_calls=0, call_reset_ts=None) + + assert policy._bucket.count() == 10 + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + frozen.tick(61) + policy.try_acquire("call", weight=1) + + def test_update_available_calls_zero_when_bucket_full_is_noop(self): + with freeze_time("2024-01-01 00:00:00"): + policy = MovingWindowCallRatePolicy(rates=[Rate(3, timedelta(minutes=1))], matchers=[]) + for _ in range(3): + policy.try_acquire("call", weight=1) + + policy.update(available_calls=0, call_reset_ts=None) + + assert policy._bucket.count() == 3 + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_available_calls_zero_ignores_expired_entries(self): + with freeze_time("2024-01-01 00:00:00") as frozen: + policy = MovingWindowCallRatePolicy( + rates=[ + Rate(5, timedelta(seconds=10)), + Rate(10, timedelta(minutes=1)), + ], + matchers=[], + ) + for _ in range(5): + policy.try_acquire("call", weight=1) + + frozen.tick(11) + policy.update(available_calls=0, call_reset_ts=None) + + assert policy._bucket.count() == 10 + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_available_calls_zero_multiple_rates(self): + with freeze_time("2024-01-01 00:00:00"): + policy = MovingWindowCallRatePolicy( + rates=[ + Rate(10, timedelta(minutes=10)), + Rate(3, timedelta(seconds=10)), + Rate(2, timedelta(hours=1)), + ], + matchers=[], + ) + policy.try_acquire("call", weight=1) + + policy.update(available_calls=0, call_reset_ts=None) + + assert policy._bucket.count() == 2 + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_positive_available_calls_is_noop(self): + with freeze_time("2024-01-01 00:00:00"): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + for _ in range(3): + policy.try_acquire("call", weight=1) + + policy.update(available_calls=5, call_reset_ts=None) + assert policy._bucket.count() == 3 + + policy.update(available_calls=0, call_reset_ts=datetime.now()) + assert policy._bucket.count() == 3 + class TestHttpStreamIntegration: def test_without_cache(self, mocker, requests_mock):