diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index aa3bbbc5e..8482c53cd 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2014,6 +2014,11 @@ definitions: type: array items: "$ref": "#/definitions/HttpRequestRegexMatcher" + max_header_driven_wait: + title: Maximum Header-Driven Wait + description: Maximum wait a rate-limit-header-driven update may induce. Rates with longer windows are not adjusted from headers. Defaults to PT10M. + type: string + examples: ["PT10M", "PT1M"] additionalProperties: true UnlimitedCallRatePolicy: title: Unlimited Call Rate Policy diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 21a92cc3b..df6721723 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -2217,6 +2217,12 @@ class Config: description="List of matchers that define which requests this policy applies to.", title="Matchers", ) + max_header_driven_wait: Optional[str] = Field( + None, + description="Maximum wait a rate-limit-header-driven update may induce. Rates with longer windows are not adjusted from headers. Defaults to PT10M.", + examples=["PT10M", "PT1M"], + title="Maximum Header-Driven Wait", + ) class UnlimitedCallRatePolicy(BaseModel): diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index aa546c73a..17e7775ac 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4603,6 +4603,11 @@ def create_moving_window_call_rate_policy( return MovingWindowCallRatePolicy( rates=rates, matchers=matchers, + **( + {"max_header_driven_wait": parse_duration(model.max_header_driven_wait)} + if model.max_header_driven_wait + else {} + ), ) def create_unlimited_call_rate_policy( diff --git a/airbyte_cdk/sources/streams/call_rate.py b/airbyte_cdk/sources/streams/call_rate.py index 4a06db3b2..0636e1491 100644 --- a/airbyte_cdk/sources/streams/call_rate.py +++ b/airbyte_cdk/sources/streams/call_rate.py @@ -15,7 +15,7 @@ import requests import requests_cache -from pyrate_limiter import InMemoryBucket, Limiter, RateItem, TimeClock +from pyrate_limiter import InMemoryBucket, Limiter, RateItem, TimeClock, binary_search from pyrate_limiter import Rate as PyRateRate from pyrate_limiter.exceptions import BucketFullException @@ -429,19 +429,35 @@ class MovingWindowCallRatePolicy(BaseCallRatePolicy): This strategy requires saving of timestamps of all requests within a window. """ - def __init__(self, rates: list[Rate], matchers: list[RequestMatcher]): + # Header-driven updates must not induce waits longer than this; sources heartbeat well above it. + # Policies with longer windows are left unchanged rather than parking a worker for the whole window. + DEFAULT_MAX_HEADER_DRIVEN_WAIT = timedelta(minutes=10) + + def __init__( + self, + rates: list[Rate], + matchers: list[RequestMatcher], + max_header_driven_wait: timedelta = DEFAULT_MAX_HEADER_DRIVEN_WAIT, + ): """Constructor :param rates: list of rates, the order is important and must be ascending :param matchers: + :param max_header_driven_wait: maximum wait a header-driven update may induce. Rates whose + windows exceed this value are excluded from header-driven updates, so a policy with no + rate inside the bound is never adjusted from response headers. Defaults to 10 minutes, + chosen to stay well inside the platform's source heartbeat. """ if not rates: raise ValueError("The list of rates can not be empty") + if max_header_driven_wait <= timedelta(0): + raise ValueError("max_header_driven_wait must be positive") pyrate_rates = [ PyRateRate(limit=rate.limit, interval=int(rate.interval.total_seconds() * 1000)) for rate in rates ] self._bucket = InMemoryBucket(pyrate_rates) + self._max_header_driven_wait_ms = int(max_header_driven_wait.total_seconds() * 1000) # Limiter will create the background task that clears old requests in the bucket self._limiter = Limiter(self._bucket) super().__init__(matchers=matchers) @@ -478,22 +494,64 @@ def update( ) -> None: """Adjust call bucket to reflect the state of the API server - :param available_calls: - :param call_reset_ts: + The bucket is filled with dummy calls until what it still allows matches what the API + reports as available. Updates only ever lower the local allowance: when the API reports + more available calls than the configured rates allow, the configured rates win. + + `call_reset_ts` is not used. A moving window has no reset point, so the only actionable + part of the API feedback is the number of calls left; the window length stays the one + the rates were configured with. + + When several rates are configured, the update applies to the most constraining one. A + header describing a coarser window only starts to bite near the end of that window; + mapping headers to a specific rate is deliberately out of scope. + + The subtraction below compares `_calls_left`, which counts bucket entries (weight units; + `put` stores `weight` copies), with `available_calls`, which counts requests. They + coincide for unweighted policies but diverge when weighted matchers are used with a + remaining header. + + :param available_calls: number of calls the API reports as still available + :param call_reset_ts: unused, see above :return: """ - if ( - available_calls is not None and call_reset_ts is None - ): # 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 - 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)) - # 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() + if available_calls is None: + return + + available_calls = max(0, available_calls) + with self._limiter.lock: + now: int = TimeClock().now() # type: ignore[no-untyped-call] + calls_left = self._calls_left(now) + if calls_left is None: + return + + items_to_add = calls_left - available_calls + if items_to_add > 0: + logger.debug( + "got rate limit update from api, adjusting available calls from %s to %s", + calls_left, + available_calls, + ) + if not self._bucket.put(RateItem(name="dummy", timestamp=now, weight=items_to_add)): + logger.warning( + "could not adjust available calls: calls_left=%s, available_calls=%s, " + "rejected_weight=%s", + calls_left, + available_calls, + items_to_add, + ) + + def _calls_left(self, now: int) -> Optional[int]: + """Number of calls the bucket still allows, i.e. the most constraining of all rates.""" + items = self._bucket.items + calls_left = [] + for rate in self._bucket.rates: + if rate.interval > self._max_header_driven_wait_ms: + continue + lower_bound_idx = binary_search(items, now - rate.interval) + calls_used = len(items) - lower_bound_idx if lower_bound_idx >= 0 else 0 + calls_left.append(rate.limit - calls_used) + return min(calls_left) if calls_left else None def __str__(self) -> str: """Return a human-friendly description of the moving window rate policy for logging purposes.""" diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index 21c99adc7..77af36b4d 100644 --- a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +++ b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py @@ -4879,6 +4879,7 @@ def test_api_budget(): "interval": "PT0.1S", # 0.1 seconds } ], + "max_header_driven_wait": "PT2M", "matchers": [ { "type": "HttpRequestRegexMatcher", @@ -4938,6 +4939,7 @@ def test_api_budget(): # The 0.1s from 'PT0.1S' is stored in ms by PyRateLimiter internally # but here just check that the limit and interval exist assert policy._bucket.rates[0].interval == 100 # 100 ms + assert policy._max_header_driven_wait_ms == 120_000 def test_api_budget_passed_to_custom_requester(): @@ -4981,6 +4983,9 @@ def test_api_budget_passed_to_custom_requester(): assert isinstance(custom_requester.api_budget, HttpAPIBudget) assert custom_requester._http_client._api_budget is custom_requester.api_budget assert len(custom_requester._http_client._api_budget._policies) == 1 + policy = custom_requester.api_budget._policies[0] + assert isinstance(policy, MovingWindowCallRatePolicy) + assert policy._max_header_driven_wait_ms == 600_000 def test_api_budget_does_not_override_custom_requester_default_value(): diff --git a/unit_tests/sources/streams/test_call_rate.py b/unit_tests/sources/streams/test_call_rate.py index a423fe573..4c940b7cb 100644 --- a/unit_tests/sources/streams/test_call_rate.py +++ b/unit_tests/sources/streams/test_call_rate.py @@ -16,6 +16,7 @@ APIBudget, CallRateLimitHit, FixedWindowCallRatePolicy, + HttpAPIBudget, HttpRequestMatcher, HttpRequestRegexMatcher, MovingWindowCallRatePolicy, @@ -256,6 +257,15 @@ def test_no_rates(self): with pytest.raises(ValueError, match="The list of rates can not be empty"): MovingWindowCallRatePolicy(rates=[], matchers=[]) + @pytest.mark.parametrize("max_header_driven_wait", [timedelta(0), timedelta(minutes=-1)]) + def test_invalid_max_header_driven_wait(self, max_header_driven_wait): + with pytest.raises(ValueError, match="max_header_driven_wait must be positive"): + MovingWindowCallRatePolicy( + rates=[Rate(10, timedelta(minutes=1))], + matchers=[], + max_header_driven_wait=max_header_driven_wait, + ) + 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=[]) @@ -305,6 +315,170 @@ 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_with_reset_ts(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=2, call_reset_ts=datetime.now()) + + policy.try_acquire("call", weight=1) + policy.try_acquire("call", weight=1) + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_only_lowers_allowance(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=50, call_reset_ts=datetime.now()) + + for _ in range(10): + policy.try_acquire("call", weight=1) + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_is_noop_without_available_calls(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=None, call_reset_ts=datetime.now()) + policy.update(available_calls=None, call_reset_ts=None) + + for _ in range(10): + policy.try_acquire("call", weight=1) + + def test_update_respects_the_most_constraining_rate(self): + policy = MovingWindowCallRatePolicy( + rates=[ + Rate(10, timedelta(seconds=1)), + Rate(5, timedelta(minutes=1)), + ], + matchers=[], + ) + + policy.update(available_calls=1, call_reset_ts=None) + + policy.try_acquire("call", weight=1) + with pytest.raises(CallRateLimitHit) as exc: + policy.try_acquire("call", weight=1) + assert exc.value.rate == "limit=5/1.0m" + assert exc.value.time_to_wait.total_seconds() == pytest.approx(60, 0.1) + + def test_update_available_calls_zero_fills_bucket(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=0, call_reset_ts=None) + + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_ignores_rates_over_header_wait_cap(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(100, timedelta(minutes=15))], matchers=[]) + + policy.update(available_calls=0, call_reset_ts=None) + + for _ in range(100): + policy.try_acquire("call", weight=1) + + def test_update_uses_configured_header_wait_cap(self): + policy = MovingWindowCallRatePolicy( + rates=[Rate(100, timedelta(minutes=15))], + matchers=[], + max_header_driven_wait=timedelta(minutes=20), + ) + + policy.update(available_calls=0, call_reset_ts=None) + + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_uses_tightened_header_wait_cap(self): + default_policy = MovingWindowCallRatePolicy( + rates=[Rate(100, timedelta(minutes=5))], matchers=[] + ) + tightened_policy = MovingWindowCallRatePolicy( + rates=[Rate(100, timedelta(minutes=5))], + matchers=[], + max_header_driven_wait=timedelta(minutes=1), + ) + + default_policy.update(available_calls=0, call_reset_ts=None) + tightened_policy.update(available_calls=0, call_reset_ts=None) + + with pytest.raises(CallRateLimitHit): + default_policy.try_acquire("call", weight=1) + for _ in range(100): + tightened_policy.try_acquire("call", weight=1) + + def test_update_caps_to_eligible_rate(self): + policy = MovingWindowCallRatePolicy( + rates=[ + Rate(10, timedelta(minutes=1)), + Rate(100, timedelta(minutes=15)), + ], + matchers=[], + ) + + policy.update(available_calls=0, call_reset_ts=None) + + with pytest.raises(CallRateLimitHit) as exc: + policy.try_acquire("call", weight=1) + assert exc.value.time_to_wait.total_seconds() <= 600 + + def test_update_clamps_negative_available_calls(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=-1, call_reset_ts=None) + + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + +class TestHttpAPIBudget: + def test_update_from_response(self, mocker): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + budget = HttpAPIBudget(policies=[policy]) + request = Request("GET", "https://example.com") + response = mocker.Mock(spec=requests.Response) + response.headers = requests.structures.CaseInsensitiveDict( + { + "RateLimit-Remaining": "1", + "RateLimit-Reset": "60", + "RateLimit-Limit": "60", + } + ) + response.status_code = 200 + + budget.update_from_response(request, response) + + budget.acquire_call(request, block=False) + with pytest.raises(CallRateLimitHit): + budget.acquire_call(request, block=False) + + def test_update_from_429_response(self, mocker): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + budget = HttpAPIBudget(policies=[policy]) + request = Request("GET", "https://example.com") + response = mocker.Mock(spec=requests.Response) + response.headers = requests.structures.CaseInsensitiveDict() + response.status_code = 429 + + budget.update_from_response(request, response) + + with pytest.raises(CallRateLimitHit) as exc: + budget.acquire_call(request, block=False) + assert exc.value.time_to_wait.total_seconds() <= 600 + + def test_update_from_429_response_ignores_over_cap_policy(self, mocker): + policy = MovingWindowCallRatePolicy(rates=[Rate(100, timedelta(minutes=15))], matchers=[]) + budget = HttpAPIBudget(policies=[policy]) + request = Request("GET", "https://example.com") + response = mocker.Mock(spec=requests.Response) + response.headers = requests.structures.CaseInsensitiveDict() + response.status_code = 429 + + budget.update_from_response(request, response) + + for _ in range(100): + budget.acquire_call(request, block=False) + class TestHttpStreamIntegration: def test_without_cache(self, mocker, requests_mock):