From d3607dceb6587b622024f912087f57a27264f7f5 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:48:30 +0000 Subject: [PATCH] fix: enforce total timeout when acquiring API budget call credit Co-Authored-By: bot_apk --- airbyte_cdk/sources/streams/call_rate.py | 35 +++--- unit_tests/sources/streams/test_call_rate.py | 112 +++++++++++++++++++ 2 files changed, 133 insertions(+), 14 deletions(-) diff --git a/airbyte_cdk/sources/streams/call_rate.py b/airbyte_cdk/sources/streams/call_rate.py index 4a06db3b22..70b6c8b6bc 100644 --- a/airbyte_cdk/sources/streams/call_rate.py +++ b/airbyte_cdk/sources/streams/call_rate.py @@ -622,11 +622,12 @@ def _do_acquire( :param request: the API request :param policy: the matching rate-limiting policy :param block: indicates whether to block until a call credit is available - :param timeout: maximum time to wait if blocking + :param timeout: maximum total time to wait if blocking, in seconds :raises: CallRateLimitHit if unable to acquire a call credit """ last_exception = None endpoint = self._extract_endpoint(request) + deadline = time.monotonic() + timeout if (block and timeout is not None) else None # 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): try: @@ -635,25 +636,31 @@ def _do_acquire( return except CallRateLimitHit as exc: last_exception = exc - if block: - if timeout is not None: - time_to_wait = min(timedelta(seconds=timeout), exc.time_to_wait) - else: - time_to_wait = exc.time_to_wait - # Ensure we never sleep for a negative duration. - time_to_wait = max(timedelta(0), time_to_wait) - logger.debug( - f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}). " - f"Sleeping for {time_to_wait} on attempt {attempt}." - ) - time.sleep(time_to_wait.total_seconds()) - else: + if not block: logger.debug( f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}) " f"and blocking is disabled." ) raise + time_to_wait = exc.time_to_wait + if deadline is not None: + remaining = timedelta(seconds=deadline - time.monotonic()) + if remaining <= timedelta(0): + logger.debug( + f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}) " + f"and the total timeout of {timeout}s has been exhausted." + ) + raise + time_to_wait = min(time_to_wait, remaining) + # Ensure we never sleep for a negative duration. + time_to_wait = max(timedelta(0), time_to_wait) + logger.debug( + f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}). " + f"Sleeping for {time_to_wait} on attempt {attempt}." + ) + time.sleep(time_to_wait.total_seconds()) + if last_exception: logger.debug( f"Exhausted all {self._maximum_attempts_to_acquire} attempts to acquire a call for endpoint {endpoint} " diff --git a/unit_tests/sources/streams/test_call_rate.py b/unit_tests/sources/streams/test_call_rate.py index a423fe5737..e98a63d6a5 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, @@ -60,6 +61,117 @@ def enable_cache_fixture(): os.environ[ENV_REQUEST_CACHE_PATH] = prev_cache_path +class StubCallRatePolicy(AbstractCallRatePolicy): + def __init__(self, failures: Optional[int] = None): + self._failures = failures + self.try_acquire_calls = 0 + + def matches(self, request: Any) -> bool: + return True + + def try_acquire(self, request: Any, weight: int) -> None: + self.try_acquire_calls += 1 + if self._failures is None or self._failures > 0: + if self._failures is not None: + self._failures -= 1 + raise CallRateLimitHit( + error="x", + item=request, + weight=1, + rate="r", + time_to_wait=timedelta(seconds=60), + ) + + def update(self, available_calls: Optional[int], call_reset_ts: Optional[datetime]) -> None: + pass + + +@pytest.fixture +def fake_clock(mocker): + class FakeClock: + def __init__(self): + self.now = 0.0 + self.sleep_calls = [] + + def monotonic(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.sleep_calls.append(seconds) + self.now += seconds + + clock = FakeClock() + mocker.patch( + "airbyte_cdk.sources.streams.call_rate.time.monotonic", + side_effect=clock.monotonic, + ) + mocker.patch( + "airbyte_cdk.sources.streams.call_rate.time.sleep", + side_effect=clock.sleep, + ) + return clock + + +class TestAPIBudgetAcquireTimeout: + def test_timeout_is_total_budget(self, fake_clock): + policy = StubCallRatePolicy() + budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4) + + with pytest.raises(CallRateLimitHit): + budget.acquire_call(object(), block=True, timeout=1) + + assert fake_clock.sleep_calls == [1.0] + assert policy.try_acquire_calls == 2 + + def test_zero_timeout_does_not_sleep(self, fake_clock): + policy = StubCallRatePolicy() + budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4) + + with pytest.raises(CallRateLimitHit): + budget.acquire_call(object(), block=True, timeout=0) + + assert fake_clock.sleep_calls == [] + assert policy.try_acquire_calls == 1 + + def test_succeeds_before_timeout(self, fake_clock): + policy = StubCallRatePolicy(failures=2) + budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4) + + budget.acquire_call(object(), block=True, timeout=200) + + assert fake_clock.sleep_calls == [60.0, 60.0] + assert policy.try_acquire_calls == 3 + + def test_none_timeout_preserves_success_behavior(self, fake_clock): + policy = StubCallRatePolicy(failures=2) + budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4) + + budget.acquire_call(object(), block=True, timeout=None) + + assert fake_clock.sleep_calls == [60.0, 60.0] + assert policy.try_acquire_calls == 3 + + def test_none_timeout_preserves_attempt_behavior(self, fake_clock): + policy = StubCallRatePolicy() + budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4) + + with pytest.raises(CallRateLimitHit): + budget.acquire_call(object(), block=True, timeout=None) + + assert fake_clock.sleep_calls == [60.0, 60.0, 60.0] + assert policy.try_acquire_calls == 3 + + def test_non_blocking_raises_immediately(self, fake_clock): + policy = StubCallRatePolicy() + budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4) + + with pytest.raises(CallRateLimitHit): + budget.acquire_call(object(), block=False, timeout=1) + + assert fake_clock.sleep_calls == [] + assert policy.try_acquire_calls == 1 + + class TestHttpRequestMatcher: try_all_types_of_requests = pytest.mark.parametrize( "request_factory",