From 02320389faae032a2c2196d1ecb27332d6e5dda3 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:50:55 +0000 Subject: [PATCH] fix: apply redirect response rate-limit headers to the request that produced them Co-Authored-By: bot_apk --- airbyte_cdk/sources/streams/call_rate.py | 8 +- unit_tests/sources/streams/test_call_rate.py | 114 +++++++++++++++++++ 2 files changed, 121 insertions(+), 1 deletion(-) diff --git a/airbyte_cdk/sources/streams/call_rate.py b/airbyte_cdk/sources/streams/call_rate.py index 4a06db3b22..b049575f0f 100644 --- a/airbyte_cdk/sources/streams/call_rate.py +++ b/airbyte_cdk/sources/streams/call_rate.py @@ -727,7 +727,13 @@ def send(self, request: requests.PreparedRequest, **kwargs: Any) -> requests.Res """Send a request with rate-limiting.""" self._api_budget.acquire_call(request) response = super().send(request, **kwargs) - self._api_budget.update_from_response(request, response) + # When redirects are followed, each hop is sent through this method with + # ``allow_redirects=False`` and updates the budget for its own request/response. + # The outer call must therefore pair ``request`` with the response it actually + # produced (the first one), not with the final response of the redirect chain. + response_history = getattr(response, "history", []) + own_response = response_history[0] if response_history else response + self._api_budget.update_from_response(request, own_response) return response diff --git a/unit_tests/sources/streams/test_call_rate.py b/unit_tests/sources/streams/test_call_rate.py index a423fe5737..1236d5b71a 100644 --- a/unit_tests/sources/streams/test_call_rate.py +++ b/unit_tests/sources/streams/test_call_rate.py @@ -16,8 +16,10 @@ APIBudget, CallRateLimitHit, FixedWindowCallRatePolicy, + HttpAPIBudget, HttpRequestMatcher, HttpRequestRegexMatcher, + LimiterSession, MovingWindowCallRatePolicy, Rate, UnlimitedCallRatePolicy, @@ -477,6 +479,118 @@ def test_moving_window_rejects_weight_exceeding_limit(self): policy.try_acquire(req, weight=50) +class TestLimiterSessionRedirects: + @staticmethod + def _build_budget(policies): + return HttpAPIBudget( + ratelimit_remaining_header="ratelimit-remaining", + ratelimit_reset_header="ratelimit-reset", + policies=policies, + ) + + @staticmethod + def _policy(matchers): + return FixedWindowCallRatePolicy( + next_reset_ts=datetime.now() + timedelta(minutes=1), + period=timedelta(minutes=1), + call_limit=100, + matchers=matchers, + ) + + @staticmethod + def _response(request, status_code, headers): + response = requests.Response() + response.status_code = status_code + response.headers.update(headers) + response.url = request.url + response.request = request + response._content = b"" + return response + + def test_redirect_updates_only_target_policy(self, mocker): + responses = [ + (302, {"Location": "https://api.example.com/b", "ratelimit-remaining": "99"}), + (200, {"ratelimit-remaining": "0"}), + ] + + def fake_send(_adapter, request, **_kwargs): + status_code, headers = responses.pop(0) + return self._response(request, status_code, headers) + + adapter_send = mocker.patch.object( + requests.adapters.HTTPAdapter, "send", autospec=True, side_effect=fake_send + ) + budget = self._build_budget( + [ + self._policy([HttpRequestMatcher(url="https://api.example.com/a")]), + self._policy([HttpRequestMatcher(url="https://api.example.com/b")]), + ] + ) + + LimiterSession(api_budget=budget).get("https://api.example.com/a") + + with pytest.raises(CallRateLimitHit): + budget.acquire_call(Request("GET", "https://api.example.com/b").prepare(), block=False) + budget.acquire_call(Request("GET", "https://api.example.com/a").prepare(), block=False) + assert adapter_send.call_count == 2 + + def test_redirect_with_shared_policy(self, mocker): + responses = [ + (302, {"Location": "https://api.example.com/b", "ratelimit-remaining": "50"}), + (200, {"ratelimit-remaining": "0"}), + ] + + def fake_send(_adapter, request, **_kwargs): + status_code, headers = responses.pop(0) + return self._response(request, status_code, headers) + + mocker.patch.object( + requests.adapters.HTTPAdapter, "send", autospec=True, side_effect=fake_send + ) + matcher = HttpRequestRegexMatcher( + url_base="https://api.example.com", url_path_pattern=r"^/(a|b)$" + ) + budget = self._build_budget([self._policy([matcher])]) + + LimiterSession(api_budget=budget).get("https://api.example.com/a") + + with pytest.raises(CallRateLimitHit): + budget.acquire_call(Request("GET", "https://api.example.com/a").prepare(), block=False) + with pytest.raises(CallRateLimitHit): + budget.acquire_call(Request("GET", "https://api.example.com/b").prepare(), block=False) + + responses.extend( + [ + (302, {"Location": "https://api.example.com/b", "ratelimit-remaining": "50"}), + (200, {"ratelimit-remaining": "10"}), + ] + ) + budget = self._build_budget([self._policy([matcher])]) + LimiterSession(api_budget=budget).get("https://api.example.com/a") + budget.acquire_call(Request("GET", "https://api.example.com/b").prepare(), block=False) + + def test_no_redirect_unchanged(self, mocker): + def fake_send(_adapter, request, **_kwargs): + return self._response(request, 200, {"ratelimit-remaining": "0"}) + + adapter_send = mocker.patch.object( + requests.adapters.HTTPAdapter, "send", autospec=True, side_effect=fake_send + ) + budget = self._build_budget( + [ + self._policy([HttpRequestMatcher(url="https://api.example.com/a")]), + self._policy([HttpRequestMatcher(url="https://api.example.com/b")]), + ] + ) + + LimiterSession(api_budget=budget).get("https://api.example.com/a") + + with pytest.raises(CallRateLimitHit): + budget.acquire_call(Request("GET", "https://api.example.com/a").prepare(), block=False) + budget.acquire_call(Request("GET", "https://api.example.com/b").prepare(), block=False) + assert adapter_send.call_count == 1 + + class TestHttpRequestRegexMatcher: """ Tests for the new regex-based logic: