From 49db636cb6ac94e57fdf0569727c1f9efebce88d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 14 Sep 2026 23:53:59 +0000 Subject: [PATCH] fix: block MovingWindowCallRatePolicy until server reset when quota is exhausted Co-Authored-By: bot_apk --- airbyte_cdk/sources/streams/call_rate.py | 27 ++++- unit_tests/sources/streams/test_call_rate.py | 108 +++++++++++++++++++ 2 files changed, 134 insertions(+), 1 deletion(-) diff --git a/airbyte_cdk/sources/streams/call_rate.py b/airbyte_cdk/sources/streams/call_rate.py index 4a06db3b22..57b0587bf5 100644 --- a/airbyte_cdk/sources/streams/call_rate.py +++ b/airbyte_cdk/sources/streams/call_rate.py @@ -442,6 +442,7 @@ def __init__(self, rates: list[Rate], matchers: list[RequestMatcher]): for rate in rates ] self._bucket = InMemoryBucket(pyrate_rates) + self._reset_ts: Optional[datetime.datetime] = None # Limiter will create the background task that clears old requests in the bucket self._limiter = Limiter(self._bucket) super().__init__(matchers=matchers) @@ -455,6 +456,19 @@ def try_acquire(self, request: Any, weight: int) -> None: f"Weight can not exceed the lowest configured rate limit ({lowest_limit})" ) + with self._limiter.lock: + if self._reset_ts is not None: + now = datetime.datetime.now() + if now < self._reset_ts: + raise CallRateLimitHit( + error=f"API call budget exhausted until {self._reset_ts.isoformat()}", + item=request, + weight=weight, + rate="api-reset", + time_to_wait=self._reset_ts - now, + ) + self._reset_ts = None + try: self._limiter.try_acquire(request, weight=weight) except BucketFullException as exc: @@ -476,12 +490,22 @@ def try_acquire(self, request: Any, weight: int) -> None: def update( self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime] ) -> None: - """Adjust call bucket to reflect the state of the API server + """Adjust the call bucket to reflect the state of the API server. :param available_calls: :param call_reset_ts: :return: """ + if available_calls == 0 and call_reset_ts is not None: + with self._limiter.lock: + if call_reset_ts > datetime.datetime.now(): + logger.debug( + "got rate limit update from api, no calls available until %s", + call_reset_ts, + ) + self._reset_ts = call_reset_ts + return + if ( available_calls is not None and call_reset_ts is None ): # we do our best to sync buckets with API @@ -505,6 +529,7 @@ def __str__(self) -> str: matcher_str = ", ".join(f"{matcher}" for matcher in self._matchers) return ( f"MovingWindowCallRatePolicy(rates=[{rates_info}], current_bucket_count={current_bucket_count}, " + f"reset_ts={self._reset_ts}, " f"matchers=[{matcher_str}])" ) diff --git a/unit_tests/sources/streams/test_call_rate.py b/unit_tests/sources/streams/test_call_rate.py index a423fe5737..ec82411a4c 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 @@ -16,6 +17,7 @@ APIBudget, CallRateLimitHit, FixedWindowCallRatePolicy, + HttpAPIBudget, HttpRequestMatcher, HttpRequestRegexMatcher, MovingWindowCallRatePolicy, @@ -256,6 +258,82 @@ def test_no_rates(self): with pytest.raises(ValueError, match="The list of rates can not be empty"): MovingWindowCallRatePolicy(rates=[], matchers=[]) + def test_update_with_reset_ts_blocks_until_reset(self): + with freeze_time("2024-01-01 00:00:00") as frozen: + policy = MovingWindowCallRatePolicy( + rates=[Rate(100, timedelta(minutes=1))], matchers=[] + ) + policy.try_acquire("call", weight=1) + policy.update( + available_calls=0, + call_reset_ts=datetime.now() + timedelta(seconds=60), + ) + + with pytest.raises(CallRateLimitHit) as excinfo: + policy.try_acquire("call", weight=1) + assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(60, 0.1) + assert policy._bucket.count() == 1 + + frozen.tick(30) + with pytest.raises(CallRateLimitHit) as excinfo: + policy.try_acquire("call", weight=1) + assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(30, 0.1) + + frozen.tick(31) + policy.try_acquire("call", weight=1) + assert policy._reset_ts is None + + def test_update_with_past_reset_ts_is_noop(self): + with freeze_time("2024-01-01 00:00:00"): + policy = MovingWindowCallRatePolicy( + rates=[Rate(100, timedelta(minutes=1))], matchers=[] + ) + policy.update( + available_calls=0, + call_reset_ts=datetime.now() - timedelta(seconds=1), + ) + + policy.try_acquire("call", weight=1) + assert policy._reset_ts is None + + def test_update_with_reset_ts_still_enforces_local_rate(self): + with freeze_time("2024-01-01 00:00:00") as frozen: + policy = MovingWindowCallRatePolicy(rates=[Rate(2, timedelta(minutes=1))], matchers=[]) + policy.try_acquire("call", weight=1) + policy.try_acquire("call", weight=1) + policy.update( + available_calls=0, + call_reset_ts=datetime.now() + timedelta(seconds=10), + ) + + frozen.tick(11) + with pytest.raises(CallRateLimitHit) as excinfo: + policy.try_acquire("call", weight=1) + assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(49, 0.1) + + def test_update_available_calls_zero_without_reset_ts_unchanged(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=0, call_reset_ts=None) + + assert policy._reset_ts is None + assert policy._bucket.count() == 4 + + def test_update_positive_available_calls_with_reset_ts_is_noop(self): + with freeze_time("2024-01-01 00:00:00"): + policy = MovingWindowCallRatePolicy( + rates=[Rate(100, timedelta(minutes=1))], matchers=[] + ) + policy.update( + available_calls=5, + call_reset_ts=datetime.now() + timedelta(seconds=60), + ) + + policy.try_acquire("call", weight=1) + def test_limit_rate(self): """try_acquire must respect configured call rate and throw CallRateLimitHit when hit the limit.""" policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) @@ -362,6 +440,36 @@ def test_with_cache(self, mocker, requests_mock): assert MovingWindowCallRatePolicy.try_acquire.call_count == 1 + def test_http_api_budget_honors_server_reset(self): + with freeze_time("2024-01-01 00:00:00") as frozen: + url = "https://api.example.com/x" + request = Request("GET", url).prepare() + budget = HttpAPIBudget( + policies=[ + MovingWindowCallRatePolicy( + rates=[Rate(100, timedelta(minutes=1))], + matchers=[HttpRequestMatcher(url=url)], + ) + ] + ) + + budget.acquire_call(request, block=False) + + response = requests.Response() + response.status_code = 429 + response.headers = { + "ratelimit-remaining": "0", + "ratelimit-reset": str(int(time.time()) + 60), + } + budget.update_from_response(request, response) + + with pytest.raises(CallRateLimitHit) as excinfo: + budget.acquire_call(request, block=False) + assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(60, 0.1) + + frozen.tick(61) + budget.acquire_call(request, block=False) + class TestWeightBasedRateLimiting: """Tests for weight-based rate limiting where different endpoints consume different amounts from a shared budget."""