fix: block MovingWindowCallRatePolicy until server reset when quota is exhausted

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}])"
         )
 

unit_tests/sources/streams/test_call_rate.py

@@ -9,13 +9,15 @@
 
 import pytest
 import requests
+from freezegun import freeze_time
 from requests import Request
 
 from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources.streams.call_rate import (
     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."""