diff --git a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py index a07e6a36d8..f56797109a 100644 --- a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py +++ b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py @@ -28,6 +28,12 @@ class TokenQuota: `matchers` classify outgoing requests into the pool; a pool with no matchers acts as the default pool. `remaining_path`/`reset_path`/`limit_path` locate the pool's values in the quota status response. + + `remaining_header`/`reset_header`/`limit_header`/`exhaustion_status_codes` are optional and + enable reconciling the pool against what the server reports on each response. Without them + the pool is only ever seeded from `quota_status_url`, so its counters drift whenever the + token is shared with another client, requests are in flight concurrently, or the sync runs + long enough for a single seeding to go stale. """ name: str @@ -35,6 +41,19 @@ class TokenQuota: reset_path: List[str] limit_path: Optional[List[str]] = None matchers: List[RequestMatcher] = field(default_factory=list) + remaining_header: Optional[str] = None + reset_header: Optional[str] = None + limit_header: Optional[str] = None + exhaustion_status_codes: List[int] = field(default_factory=list) + + @property + def is_response_aware(self) -> bool: + return bool( + self.remaining_header + or self.reset_header + or self.limit_header + or self.exhaustion_status_codes + ) @dataclass @@ -58,14 +77,25 @@ class RateLimitedMultipleTokenAuthenticator(DeclarativeAuthenticator): (`max(budget_min_reserve, budget_reserve_fraction * limit)`), a small delay proportional to `seconds_until_reset / total_remaining` (capped at 10s) is injected before each request. + Implements `ResponseAwareAuthenticator` and `TokenRotatingAuthenticator` (see + `airbyte_cdk.sources.streams.http.requests_native_auth.protocols`), which is how `HttpClient` + feeds it responses and asks it whether a rate-limit wait can be skipped. + Counters are seeded per token from `quota_status_url` on first use and refreshed after an - exhaustion wait. All state transitions are guarded by a lock so the authenticator can be - shared safely across concurrent streams; sleeps never hold the lock. + exhaustion wait. When a pool declares response headers, `update_from_response` additionally + reconciles that pool against the server on every response, which keeps the counters honest + between seedings and makes the authenticator rotate off a token the server has rejected even + though the local count still looks healthy. All state transitions are guarded by a lock so + the authenticator can be shared safely across concurrent streams; sleeps never hold the lock. """ HEARTBEAT_INTERVAL = 60.0 # Log every 60s during exhaustion wait MAX_BUDGET_DELAY = 10.0 # Cap for the per-request proactive throttling delay MIN_EXHAUSTION_WAIT = 5.0 # Floor for the exhaustion wait, so stale reset timestamps can't cause a refresh busy-loop + # How far behind the reset we hold a response's reset may be while still counting as the + # current window. Covers ordinary disagreement between the quota endpoint and response + # headers; anything older is treated as belonging to a window that has already rolled over. + RESET_SKEW_TOLERANCE = timedelta(seconds=60) def __init__( self, @@ -319,3 +349,145 @@ def _extract_path(self, response_body: Mapping[str, Any], path: List[str]) -> An ) value = value[key] return value + + def update_from_response( + self, request: requests.PreparedRequest, response: requests.Response + ) -> None: + """Reconcile the matched pool's counters against what the server reported. + + Called once per HTTP attempt by `HttpClient`. Inert unless the matched pool declares + response headers, so behaviour is unchanged for pools configured only with paths. + + The update is attributed to the token that *sent* the request rather than to the + currently active one: under concurrency the authenticator may have rotated between the + request going out and the response coming back. + """ + quota = self._match_quota(request) + if not quota.is_response_aware: + return + token = self._token_from_request(request) + if token is None: + return + + remaining = self._header_int(response, quota.remaining_header) + if remaining is None and response.status_code in quota.exhaustion_status_codes: + # The server rejected the call for rate-limit reasons but told us nothing about the + # remaining count. Treat the pool as spent so the next request rotates. + remaining = 0 + reset_at = self._header_datetime(response, quota.reset_header) + limit = self._header_int(response, quota.limit_header) + if remaining is None and limit is None and reset_at is None: + return + + with self._lock: + state = self._states.get(token, {}).get(quota.name) + if state is None: + return # not seeded yet; the initial seeding is the more authoritative source + if limit is not None and limit > 0 and (reset_at is None or reset_at >= state.reset_at): + # A response from an older window carries that window's limit. Taking it would + # skew the throttling reserve and, on a later reset-only response, refill the + # pool to the wrong capacity -- so the limit is accepted on the same terms as + # the reset itself. Assigned before the rollover branch below, which reads + # `state.limit` when a fresh window arrives with no remaining header. + state.limit = limit + if reset_at is not None and reset_at > state.reset_at: + # The quota window rolled over, so the local count is meaningless -- take the + # server's numbers wholesale. With no remaining header to go on, a fresh window + # is worth a full limit. + state.reset_at = reset_at + state.remaining = remaining if remaining is not None else state.limit + elif remaining is not None and ( + (remaining <= 0 and response.status_code in quota.exhaustion_status_codes) + or reset_at is None + or reset_at >= state.reset_at - self.RESET_SKEW_TOLERANCE + ): + # Same window: only ever tighten the estimate. Responses arrive out of order and + # a slow one carries a stale, higher count, so handing those calls back would let + # concurrent requests overspend. The window itself is never moved backwards. + # + # A count from a window that has already rolled over describes a window that no + # longer exists, and `min` would pin the fresh pool to it for the rest of the + # hour, so those are ignored -- with two exceptions. + # + # First, a zero on a response the pool counts as rate-limited is an exhaustion + # signal and must never be dropped, or a rate limit whose reset header trails the + # value we hold would silently stop rotation. The status check is what keeps that + # narrow: a zero on a *successful* response is just the last call of a window, and + # honouring it from a dead window would park a pool that has already refilled. + # + # Second, a reset within `RESET_SKEW_TOLERANCE` is treated as the current window, + # since the quota endpoint and the response headers can disagree by a little. + state.remaining = min(state.remaining, remaining) + + def has_alternative_token(self, request: requests.PreparedRequest) -> bool: + """Whether another token could serve this request right now. + + Answers the question a rate-limit backoff cannot answer for itself: the wait computed + from a response's reset header assumes the only way forward is for that quota to come + back, which is false when a different token still has calls. `HttpClient` uses this to + retry promptly instead of sleeping out a window it does not need. + + Deliberately narrow. It reports True only when the token that *sent* the request is + spent for the matched pool -- so the next request is guaranteed to rotate -- and some + other token is not. If the sending token still has calls locally, the rejection was not + about exhausting it (a secondary limit, say, which on many APIs is per-user and would + reject every token alike), and waiting remains the right response. + """ + quota = self._match_quota(request) + sender = self._token_from_request(request) + with self._lock: + if not self._states or sender is None: + return False + if self._states[sender][quota.name].remaining > 0: + return False + return any( + self._states[token][quota.name].remaining > 0 + for token in self._tokens + if token != sender + ) + + def _token_from_request(self, request: requests.PreparedRequest) -> Optional[str]: + """Recover the token a request was signed with from its auth header. + + The prefix is checked rather than assumed: the header may have been written by + something other than this authenticator, and slicing blindly would leave membership in + `_states` as the only thing standing between a mangled value and a wrong attribution. + """ + value = request.headers.get(self._header) + if not value: + return None + if self._auth_method: + prefix = f"{self._auth_method} " + if not value.startswith(prefix): + return None + token = value[len(prefix) :].strip() + else: + token = value.strip() + return token if token in self._states else None + + @staticmethod + def _header_int(response: requests.Response, header: Optional[str]) -> Optional[int]: + if not header: + return None + value = response.headers.get(header) + if value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + @staticmethod + def _header_datetime( + response: requests.Response, header: Optional[str] + ) -> Optional[AirbyteDateTime]: + if not header: + return None + value = response.headers.get(header) + if value is None: + return None + try: + # Same parsing rules as `reset_path`, so epoch seconds and ISO 8601 both work. + return ab_datetime_parse(value) + except Exception: + return None diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 97dda3b98b..091c694d23 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -497,6 +497,32 @@ definitions: type: array items: "$ref": "#/definitions/HttpRequestRegexMatcher" + remaining_header: + title: Remaining Header + description: Optional response header carrying the remaining call count for this pool. When set, the pool's counter is reconciled against this header on every response, which corrects drift caused by sharing the token with other clients, by requests in flight concurrently, or by a sync running long enough for the initial quota status read to go stale. Without it the pool is only ever seeded from the quota status endpoint. + type: string + examples: + - "X-RateLimit-Remaining" + reset_header: + title: Reset Header + description: Optional response header carrying the quota reset timestamp for this pool. Parsed with the same rules as `reset_path`, so epoch seconds and ISO 8601 both work. Used to tell a rolled-over quota window from the current one; a response proving the window has rolled over restores the pool to its limit. Most useful alongside `remaining_header`. + type: string + examples: + - "X-RateLimit-Reset" + limit_header: + title: Limit Header + description: Optional response header carrying the total call limit for this pool, used to keep the proactive throttling reserve accurate as the limit changes. + type: string + examples: + - "X-RateLimit-Limit" + exhaustion_status_codes: + title: Exhaustion Status Codes + description: Response status codes that mean this token's pool is spent. These have two effects. A response carrying one of them but no remaining count sets the pool to zero, so the next request rotates to another token instead of waiting out the reset window. They also mark which responses may report a zero for a quota window that has already elapsed, so a rate limit whose reset header trails the value being held still stops the token being used; a zero on any other response is treated as the last call of a finished window and ignored. Leaving this empty means such trailing rejections are ignored unless their reset is within the skew tolerance of the current window. Only list codes the API uses exclusively for rate limiting -- a code that also signals other failures would park a healthy token. + type: array + items: + type: integer + examples: + - [429] $parameters: type: object additionalProperties: true diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index a84412eaf9..37eb50b7b1 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -609,6 +609,30 @@ class TokenQuota(BaseModel): description="List of matchers that classify outgoing requests into this quota pool. The first pool whose matcher matches a request is used. A pool with no matchers acts as the default pool.", title="Matchers", ) + remaining_header: Optional[str] = Field( + None, + description="Optional response header carrying the remaining call count for this pool. When set, the pool's counter is reconciled against this header on every response, which corrects drift caused by sharing the token with other clients, by requests in flight concurrently, or by a sync running long enough for the initial quota status read to go stale. Without it the pool is only ever seeded from the quota status endpoint.", + examples=["X-RateLimit-Remaining"], + title="Remaining Header", + ) + reset_header: Optional[str] = Field( + None, + description="Optional response header carrying the quota reset timestamp for this pool. Parsed with the same rules as `reset_path`, so epoch seconds and ISO 8601 both work. Used to tell a rolled-over quota window from the current one; a response proving the window has rolled over restores the pool to its limit. Most useful alongside `remaining_header`.", + examples=["X-RateLimit-Reset"], + title="Reset Header", + ) + limit_header: Optional[str] = Field( + None, + description="Optional response header carrying the total call limit for this pool, used to keep the proactive throttling reserve accurate as the limit changes.", + examples=["X-RateLimit-Limit"], + title="Limit Header", + ) + exhaustion_status_codes: Optional[List[int]] = Field( + None, + description="Response status codes that mean this token's pool is spent. These have two effects. A response carrying one of them but no remaining count sets the pool to zero, so the next request rotates to another token instead of waiting out the reset window. They also mark which responses may report a zero for a quota window that has already elapsed, so a rate limit whose reset header trails the value being held still stops the token being used; a zero on any other response is treated as the last call of a finished window and ignored. Leaving this empty means such trailing rejections are ignored unless their reset is within the skew tolerance of the current window. Only list codes the API uses exclusively for rate limiting -- a code that also signals other failures would park a healthy token.", + examples=[[429]], + title="Exhaustion Status Codes", + ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") 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 b992b4201a..a96c406d8c 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4660,6 +4660,12 @@ def create_rate_limited_multiple_token_authenticator( "remaining_path": quota_model.remaining_path, "reset_path": quota_model.reset_path, "limit_path": quota_model.limit_path, + "remaining_header": quota_model.remaining_header, + "reset_header": quota_model.reset_header, + "limit_header": quota_model.limit_header, + # Normalized the same way as the runtime TokenQuota below, so an omitted field + # and an explicit `[]` key identically and keep sharing one set of counters. + "exhaustion_status_codes": quota_model.exhaustion_status_codes or [], "matchers": [ { "method": matcher_model.method, @@ -4734,6 +4740,10 @@ def create_rate_limited_multiple_token_authenticator( remaining_path=quota_model.remaining_path, reset_path=quota_model.reset_path, limit_path=quota_model.limit_path, + remaining_header=quota_model.remaining_header, + reset_header=quota_model.reset_header, + limit_header=quota_model.limit_header, + exhaustion_status_codes=quota_model.exhaustion_status_codes or [], matchers=[ self.create_http_request_matcher(matcher_model, config) for matcher_model in quota_model.matchers or [] diff --git a/airbyte_cdk/sources/streams/http/http_client.py b/airbyte_cdk/sources/streams/http/http_client.py index c1d0eabd67..263a4ea1ba 100644 --- a/airbyte_cdk/sources/streams/http/http_client.py +++ b/airbyte_cdk/sources/streams/http/http_client.py @@ -50,6 +50,14 @@ rate_limit_default_backoff_handler, user_defined_backoff_handler, ) + +# Imported from the leaf module rather than the package: `protocols` pulls in nothing from the +# CDK, so this import cannot cycle no matter what else lands in `requests_native_auth` -- an +# authenticator there needing `HttpClient` (as the declarative one does) stays safe. +from airbyte_cdk.sources.streams.http.requests_native_auth.protocols import ( + ResponseAwareAuthenticator, + TokenRotatingAuthenticator, +) from airbyte_cdk.sources.utils.types import JsonType from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH @@ -90,6 +98,10 @@ def monkey_patched_get_item(self, key): # type: ignore # this interface is a co class HttpClient: _DEFAULT_MAX_RETRY: int = 5 _DEFAULT_MAX_TIME: int = 60 * 10 + # Backoff used in place of a rate-limit wait when another credential can serve the retry. + # Kept non-zero so a misreporting authenticator degrades to a slow retry, not a hot loop + # (the retry handler adds a second on top of whatever is returned). + TOKEN_ROTATION_BACKOFF: float = 0.1 _ACTIONS_TO_RETRY_ON = { ResponseAction.RETRY, ResponseAction.RATE_LIMITED, @@ -138,6 +150,7 @@ def __init__( self._request_attempt_count: Dict[requests.PreparedRequest, int] = {} self._disable_retries = disable_retries self._message_repository = message_repository + self._authenticator_update_failed = False @property def cache_filename(self) -> str: @@ -322,6 +335,59 @@ def _send_with_retry( stream_descriptor=StreamDescriptor(name=self._name), ) + def _can_retry_on_another_token(self, request: requests.PreparedRequest) -> bool: + """Whether the authenticator can serve this request from a different credential now. + + Opted into by implementing `TokenRotatingAuthenticator`. + """ + authenticator = getattr(self._session, "auth", None) + if not isinstance(authenticator, TokenRotatingAuthenticator): + return False + try: + return bool(authenticator.has_alternative_token(request)) + except Exception: + # Falling back to the computed wait is always safe, so never fail a retry over this. + self._logger.debug( + "Authenticator failed to report credential availability", exc_info=True + ) + return False + + def _update_authenticator_from_response( + self, request: requests.PreparedRequest, response: requests.Response + ) -> None: + """Let a quota-tracking authenticator reconcile its state against the server. + + Authenticators only ever see requests, so an authenticator that tracks per-token quota + has no way to learn that the server disagrees with its local bookkeeping. This is the + feedback channel, mirroring what `LimiterMixin.send` does for the API budget. + + Opted into by implementing `ResponseAwareAuthenticator`. + """ + authenticator = getattr(self._session, "auth", None) + if not isinstance(authenticator, ResponseAwareAuthenticator): + return + if getattr(response, "from_cache", False): + # A replayed cached response carries the rate-limit headers from whenever it was + # first fetched and consumed no quota of its own. + return + try: + authenticator.update_from_response(request, response) + except Exception: + # Quota bookkeeping must never turn an otherwise fine response into a failure. Warn + # once so a persistently broken update -- which silently degrades the connector back + # to single-token behaviour -- is at least diagnosable from default-level logs. + if not self._authenticator_update_failed: + self._authenticator_update_failed = True + self._logger.warning( + "Authenticator failed to update quota state from a response; token rotation " + "may fall back to local counters only. Further occurrences log at debug.", + exc_info=True, + ) + else: + self._logger.debug( + "Authenticator failed to update quota state from response", exc_info=True + ) + def _send( self, request: requests.PreparedRequest, @@ -349,6 +415,9 @@ def _send( except requests.RequestException as e: exc = e + if response is not None: + self._update_authenticator_from_response(request, response) + error_resolution: ErrorResolution = self._error_handler.interpret_response( response if response is not None else exc ) @@ -527,6 +596,22 @@ def _handle_error_resolution( if backoff_time: user_defined_backoff_time = backoff_time break + + if ( + user_defined_backoff_time + and error_resolution.response_action == ResponseAction.RATE_LIMITED + and self._can_retry_on_another_token(request) + ): + # The backoff was derived from this response's headers, which only describe the + # credential that was rejected. Another one has quota, and the retry re-signs the + # request, so waiting out this window would idle for nothing. + self._logger.info( + f"Rate limited on the current credential; retrying in " + f"{self.TOKEN_ROTATION_BACKOFF}s with another one instead of waiting " + f"{user_defined_backoff_time:.0f}s for the rate limit to reset." + ) + user_defined_backoff_time = self.TOKEN_ROTATION_BACKOFF + error_message = ( error_resolution.error_message or f"Request to {request.url} failed with failure type {error_resolution.failure_type}, response action {error_resolution.response_action}." diff --git a/airbyte_cdk/sources/streams/http/requests_native_auth/__init__.py b/airbyte_cdk/sources/streams/http/requests_native_auth/__init__.py index 307f91f40c..5df3a0d058 100644 --- a/airbyte_cdk/sources/streams/http/requests_native_auth/__init__.py +++ b/airbyte_cdk/sources/streams/http/requests_native_auth/__init__.py @@ -3,6 +3,7 @@ # from .oauth import Oauth2Authenticator, SingleUseRefreshTokenOauth2Authenticator +from .protocols import ResponseAwareAuthenticator, TokenRotatingAuthenticator from .token import BasicHttpAuthenticator, MultipleTokenAuthenticator, TokenAuthenticator __all__ = [ @@ -11,4 +12,6 @@ "TokenAuthenticator", "MultipleTokenAuthenticator", "BasicHttpAuthenticator", + "ResponseAwareAuthenticator", + "TokenRotatingAuthenticator", ] diff --git a/airbyte_cdk/sources/streams/http/requests_native_auth/protocols.py b/airbyte_cdk/sources/streams/http/requests_native_auth/protocols.py new file mode 100644 index 0000000000..7507b370b2 --- /dev/null +++ b/airbyte_cdk/sources/streams/http/requests_native_auth/protocols.py @@ -0,0 +1,62 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +"""Optional capabilities an authenticator can offer the HTTP client. + +Authenticators are normally write-only from the client's point of view: they sign a request and +that is the end of the conversation. An authenticator managing a quota, or several +interchangeable credentials, needs more than that -- it has to learn what the server said, and +it knows things about credential availability that the retry logic cannot work out on its own. + +These protocols are how it says so. Both are optional and dispatched structurally: `HttpClient` +checks the authenticator against the protocol and skips one that does not satisfy it, so +implementing one does not require inheriting from anything or importing the client. + +Define the methods on the class or set them on the instance. Dispatch is an `isinstance` check, +and from Python 3.12 those resolve protocol members with `inspect.getattr_static`, which does +not run `__getattr__` -- so a delegating authenticator that only exposes the method dynamically +satisfies the protocol on 3.10/3.11 and is silently skipped on 3.12+. Presence is all that is +checked either way: an implementation with the wrong signature still dispatches, and then fails +at the call. +""" + +from typing import Protocol, runtime_checkable + +import requests + + +@runtime_checkable +class ResponseAwareAuthenticator(Protocol): + """An authenticator that wants to see responses, not just sign requests. + + Authenticators tracking per-token quota need a feedback channel: without one they can only + guess at the server's view of the quota and cannot tell that a token has been rejected. + `HttpClient` calls `update_from_response` once per attempt on any authenticator that + implements this, skipping replayed cache hits -- a cached response carries stale rate-limit + headers and consumed no quota. + + Implementations must not raise for a response they cannot interpret, and must be safe to + call from multiple threads. + """ + + def update_from_response( + self, request: requests.PreparedRequest, response: requests.Response + ) -> None: + pass + + +@runtime_checkable +class TokenRotatingAuthenticator(Protocol): + """An authenticator holding several interchangeable credentials. + + A rate-limit backoff is computed from the response alone, so it assumes the only way + forward is for that quota to come back. An authenticator with a spare credential knows + better. `HttpClient` asks before sleeping out a rate-limit window. + + Implementations must only answer True when retrying immediately would actually use a + different credential -- otherwise the retry hammers the same rejected one. + """ + + def has_alternative_token(self, request: requests.PreparedRequest) -> bool: + pass diff --git a/unit_tests/sources/declarative/auth/test_rate_limited_multiple_token.py b/unit_tests/sources/declarative/auth/test_rate_limited_multiple_token.py index 395b9b1727..0a45800918 100644 --- a/unit_tests/sources/declarative/auth/test_rate_limited_multiple_token.py +++ b/unit_tests/sources/declarative/auth/test_rate_limited_multiple_token.py @@ -25,6 +25,10 @@ ModelToComponentFactory, ) from airbyte_cdk.sources.streams.call_rate import HttpRequestRegexMatcher +from airbyte_cdk.sources.streams.http.requests_native_auth import ( + ResponseAwareAuthenticator, + TokenRotatingAuthenticator, +) from airbyte_cdk.utils import AirbyteTracedException QUOTA_STATUS_URL = "https://api.example.com/rate_limit" @@ -147,6 +151,40 @@ def test_waits_and_reseeds_when_all_tokens_exhausted(requests_mock): ) +def test_elapsed_window_reseeds_from_the_server_rather_than_refilling_locally(requests_mock): + """An exhausted pool whose reset has passed must be refreshed from the quota endpoint. + + Handing the allowance back locally instead would let the connector invent quota it has no + evidence for: the reset timestamp says a new window is due, but only the server knows what + is actually left in it. A pool refilled from a stale timestamp never reseeds and never + waits, so it would serve unlimited calls against a quota it stopped tracking. + """ + body = {"resources": {"core": {"remaining": 3, "reset": int(time.time()) - 10, "limit": 3}}} + quota_status = requests_mock.get(QUOTA_STATUS_URL, json=body) + authenticator = RateLimitedMultipleTokenAuthenticator( + tokens=["token_1"], + quotas=[ + TokenQuota( + name="rest", + remaining_path=["resources", "core", "remaining"], + reset_path=["resources", "core", "reset"], + limit_path=["resources", "core", "limit"], + ) + ], + quota_status_url=QUOTA_STATUS_URL, + auth_method="token", + ) + + sleeps = [] + with patch("time.sleep", side_effect=lambda seconds: sleeps.append(seconds)): + for _ in range(12): + authenticator(_prepared_request()) + + # 3 calls per window: the initial seeding plus one reseed for each exhausted window. + assert quota_status.call_count == 4 + assert sleeps == [RateLimitedMultipleTokenAuthenticator.MIN_EXHAUSTION_WAIT] * 3 + + def test_budget_throttling_delay_injected_when_all_tokens_below_reserve(requests_mock): requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body(rest_remaining=100)) authenticator = _authenticator() @@ -281,6 +319,488 @@ def fake_sleep(seconds): assert clock["now"] <= 30 +# --- reconciling quota state against response headers --------------------------------------- + + +def _response_aware_quotas(exhaustion_status_codes=(429,)): + return [ + TokenQuota( + name="rest", + remaining_path=["resources", "core", "remaining"], + reset_path=["resources", "core", "reset"], + limit_path=["resources", "core", "limit"], + remaining_header="X-RateLimit-Remaining", + reset_header="X-RateLimit-Reset", + limit_header="X-RateLimit-Limit", + exhaustion_status_codes=list(exhaustion_status_codes), + ), + TokenQuota( + name="graphql", + remaining_path=["resources", "graphql", "remaining"], + reset_path=["resources", "graphql", "reset"], + limit_path=["resources", "graphql", "limit"], + remaining_header="X-RateLimit-Remaining", + reset_header="X-RateLimit-Reset", + matchers=[HttpRequestRegexMatcher(url_path_pattern="/graphql")], + ), + ] + + +def _response_aware_authenticator( + tokens=("token_1", "token_2"), exhaustion_status_codes=(429,), **kwargs +): + """`**kwargs` reach the authenticator, matching `_authenticator` above; quota-shaping + arguments are named explicitly so a caller passing e.g. `max_wait_time` doesn't hit a + TypeError from the quota builder.""" + return RateLimitedMultipleTokenAuthenticator( + tokens=list(tokens), + quotas=_response_aware_quotas(exhaustion_status_codes=exhaustion_status_codes), + quota_status_url=QUOTA_STATUS_URL, + auth_method="token", + **kwargs, + ) + + +def _response(status_code=200, headers=None, from_cache=False): + response = requests.Response() + response.status_code = status_code + response.headers.update(headers or {}) + if from_cache: + response.from_cache = True # type: ignore[attr-defined] # mirrors requests_cache + return response + + +def test_response_header_corrects_counter_downward(requests_mock): + """The server is the source of truth: a lower remaining count wins over the local estimate.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + assert authenticator._states["token_1"]["rest"].remaining == 4999 + + authenticator.update_from_response(request, _response(headers={"X-RateLimit-Remaining": "10"})) + + assert authenticator._states["token_1"]["rest"].remaining == 10 + + +def test_response_header_does_not_ratchet_counter_up_within_a_window(requests_mock): + """A slow response carries a stale, higher count. Handing those calls back would let concurrent + requests overspend, so within one window the estimate may only tighten.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + reset_at = authenticator._states["token_1"]["rest"].reset_at + authenticator._states["token_1"]["rest"].remaining = 10 + + authenticator.update_from_response( + request, + _response( + headers={ + "X-RateLimit-Remaining": "4000", + "X-RateLimit-Reset": str(int(reset_at.timestamp())), + } + ), + ) + + assert authenticator._states["token_1"]["rest"].remaining == 10 + + +def test_later_reset_in_response_starts_a_new_window(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + state = authenticator._states["token_1"]["rest"] + state.remaining = 3 + next_window = int(state.reset_at.timestamp()) + 3600 + + authenticator.update_from_response( + request, + _response(headers={"X-RateLimit-Remaining": "5000", "X-RateLimit-Reset": str(next_window)}), + ) + + assert state.remaining == 5000 + assert int(state.reset_at.timestamp()) == next_window + + +def test_slightly_earlier_reset_still_counts_as_the_current_window(requests_mock): + """The quota endpoint and the response headers can disagree by a little, so a reset just + behind the one we hold is the same window and its count still clamps.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + state = authenticator._states["token_1"]["rest"] + state.remaining = 100 + current_window = int(state.reset_at.timestamp()) + + authenticator.update_from_response( + request, + _response( + headers={"X-RateLimit-Remaining": "7", "X-RateLimit-Reset": str(current_window - 5)} + ), + ) + + assert state.remaining == 7 + assert int(state.reset_at.timestamp()) == current_window # never moved backwards + + +def test_stale_zero_from_a_successful_response_does_not_park_a_refilled_pool(requests_mock): + """A zero on a 200 is just the last call of a window, not a rate-limit rejection. + + The exhaustion exception to the window check exists so a rate limit whose reset header + trails cannot stop rotation; honouring a zero from *any* response would let a dead + window's final count park a pool that has already refilled -- the F9 bug via the escape + hatch. The status is what tells the two apart. + """ + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + state = authenticator._states["token_1"]["rest"] + previous_window = int(state.reset_at.timestamp()) + authenticator.update_from_response( + request, + _response( + headers={ + "X-RateLimit-Remaining": "5000", + "X-RateLimit-Reset": str(previous_window + 3600), + } + ), + ) + assert state.remaining == 5000 + + stale = {"X-RateLimit-Remaining": "0", "X-RateLimit-Reset": str(previous_window)} + authenticator.update_from_response(request, _response(status_code=200, headers=stale)) + assert state.remaining == 5000, "a zero from a successful dead-window response must be ignored" + + # ...but the same zero on a rate-limited response is an exhaustion signal and still lands + authenticator.update_from_response(request, _response(status_code=429, headers=stale)) + assert state.remaining == 0 + + +def test_count_from_a_rolled_over_window_does_not_clamp_the_fresh_pool(requests_mock): + """A slow response that lands after the window turned describes a window that no longer + exists. `min` would pin the refilled pool to that dead count for the rest of the hour, and + nothing short of the next rollover would undo it.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + state = authenticator._states["token_1"]["rest"] + previous_window = int(state.reset_at.timestamp()) + + # the window turns over + authenticator.update_from_response( + request, + _response( + headers={ + "X-RateLimit-Remaining": "5000", + "X-RateLimit-Reset": str(previous_window + 3600), + } + ), + ) + assert state.remaining == 5000 + + # ...and only now does an in-flight response from the old window arrive + authenticator.update_from_response( + request, + _response( + headers={"X-RateLimit-Remaining": "3", "X-RateLimit-Reset": str(previous_window)} + ), + ) + + assert state.remaining == 5000 + assert int(state.reset_at.timestamp()) == previous_window + 3600 + + +def test_limit_from_an_older_window_is_not_applied(requests_mock): + """The limit is accepted on the same terms as the reset it arrived with. + + Taking a previous window's limit would skew the throttling reserve and, on a later + reset-only response, refill the pool to the wrong capacity. + """ + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + state = authenticator._states["token_1"]["rest"] + current_window = int(state.reset_at.timestamp()) + + authenticator.update_from_response( + request, + _response( + headers={ + "X-RateLimit-Limit": "1000", + "X-RateLimit-Reset": str(current_window - 3600), + } + ), + ) + + assert state.limit == 5000 + + +def test_limit_header_updates_the_throttling_reserve(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + + authenticator.update_from_response(request, _response(headers={"X-RateLimit-Limit": "15000"})) + + assert authenticator._states["token_1"]["rest"].limit == 15000 + + +def test_exhaustion_status_code_zeroes_the_pool_and_next_request_rotates(requests_mock): + """Regression test for the core defect. + + GitHub rejects a request for rate-limit reasons while the local counter still reads healthy. + Before response awareness the authenticator kept handing out the same token and the error + handler waited out the whole reset window with other tokens sitting idle. + """ + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator() + request = authenticator(_prepared_request()) + assert request.headers["Authorization"] == "token token_1" + assert authenticator._states["token_1"]["rest"].remaining == 4999 # locally still healthy + + authenticator.update_from_response(request, _response(status_code=429)) + + assert authenticator._states["token_1"]["rest"].remaining == 0 + retry = authenticator(_prepared_request()) + assert retry.headers["Authorization"] == "token token_2" + + +def test_exhaustion_is_not_masked_by_an_earlier_reset_header(requests_mock): + """A rate-limited response must zero the pool even when its reset header is older than the + one we hold. Dropping the whole update in that case would silently disable rotation, since + the seeding read and the response headers routinely disagree by a few seconds.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body(reset_in_seconds=3600)) + authenticator = _response_aware_authenticator() + request = authenticator(_prepared_request()) + earlier = int(authenticator._states["token_1"]["rest"].reset_at.timestamp()) - 1800 + + authenticator.update_from_response( + request, + _response( + status_code=429, + headers={"X-RateLimit-Remaining": "0", "X-RateLimit-Reset": str(earlier)}, + ), + ) + + assert authenticator._states["token_1"]["rest"].remaining == 0 + # the window itself is never moved backwards + assert int(authenticator._states["token_1"]["rest"].reset_at.timestamp()) == earlier + 1800 + assert authenticator(_prepared_request()).headers["Authorization"] == "token token_2" + + +def test_reset_only_response_restores_the_pool_on_a_new_window(requests_mock): + """A response that proves the window rolled over is actionable on its own: the previous + count belongs to a window that no longer exists.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + state = authenticator._states["token_1"]["rest"] + state.remaining = 0 + next_window = int(state.reset_at.timestamp()) + 3600 + + authenticator.update_from_response( + request, _response(headers={"X-RateLimit-Reset": str(next_window)}) + ) + + assert state.remaining == 5000 # the pool's full limit, for the window just opened + assert int(state.reset_at.timestamp()) == next_window + + +def test_rate_limited_response_is_ignored_when_no_headers_are_configured(requests_mock): + """Pools configured only with paths keep their previous behaviour.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _authenticator(tokens=("token_1", "token_2")) + request = authenticator(_prepared_request()) + + authenticator.update_from_response( + request, _response(status_code=429, headers={"X-RateLimit-Remaining": "0"}) + ) + + assert authenticator._states["token_1"]["rest"].remaining == 4999 + assert authenticator(_prepared_request()).headers["Authorization"] == "token token_1" + + +def test_response_is_attributed_to_the_token_that_sent_it(requests_mock): + """Under concurrency the active token can move on before a response lands, so the update has + to follow the request's own auth header rather than whichever token is active now.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator() + request = authenticator(_prepared_request()) + assert request.headers["Authorization"] == "token token_1" + authenticator._active_token = ( + "token_2" # the authenticator rotated while the call was in flight + ) + + authenticator.update_from_response(request, _response(headers={"X-RateLimit-Remaining": "12"})) + + assert authenticator._states["token_1"]["rest"].remaining == 12 + assert authenticator._states["token_2"]["rest"].remaining == 5000 + + +def test_cached_response_does_not_update_quota_state(requests_mock): + """A replayed cached response carries stale headers and consumed no quota. `HttpClient` filters + those out; this pins the behaviour if anything ever calls the method directly.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + + from airbyte_cdk.sources.streams.http import HttpClient + + client = HttpClient(name="test", logger=__import__("logging").getLogger("test")) + client._session.auth = authenticator + client._update_authenticator_from_response( + request, _response(headers={"X-RateLimit-Remaining": "1"}, from_cache=True) + ) + + assert authenticator._states["token_1"]["rest"].remaining == 4999 + + +def test_graphql_response_updates_only_the_graphql_pool(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request("https://api.example.com/graphql")) + + authenticator.update_from_response(request, _response(headers={"X-RateLimit-Remaining": "42"})) + + assert authenticator._states["token_1"]["graphql"].remaining == 42 + assert authenticator._states["token_1"]["rest"].remaining == 5000 + + +def test_satisfies_both_authenticator_capability_protocols(requests_mock): + """`HttpClient` dispatches structurally, but the protocols are the declared contract -- an + implementation drifting from either one should fail here rather than silently stop being + consulted at runtime.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator() + + assert isinstance(authenticator, ResponseAwareAuthenticator) + assert isinstance(authenticator, TokenRotatingAuthenticator) + + +def test_has_alternative_token_when_sender_is_spent_and_another_is_not(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator() + request = authenticator(_prepared_request()) + authenticator.update_from_response(request, _response(status_code=429)) + + assert authenticator.has_alternative_token(request) is True + + +def test_no_alternative_token_while_the_sending_token_still_has_calls(requests_mock): + """A rejection that did not exhaust the sending token is not something rotation fixes -- + a secondary limit is typically per-account and would reject every token alike, so the + computed wait must stand.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator() + request = authenticator(_prepared_request()) + + assert authenticator.has_alternative_token(request) is False + + +def test_no_alternative_token_when_every_token_is_spent(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator() + request = authenticator(_prepared_request()) + authenticator._ensure_initialized() + for token in authenticator._tokens: + authenticator._states[token]["rest"].remaining = 0 + + assert authenticator.has_alternative_token(request) is False + + +def test_no_alternative_token_with_a_single_token(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + authenticator.update_from_response(request, _response(status_code=429)) + + assert authenticator.has_alternative_token(request) is False + + +def test_alternative_token_is_scoped_to_the_matched_quota_pool(requests_mock): + """Availability is answered per pool: exhausting graphql says nothing about rest. + + The graphql pool here declares `remaining_header` but no `exhaustion_status_codes`, so it + learns it is spent from the header rather than from the status code. + """ + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator() + graphql_request = authenticator(_prepared_request("https://api.example.com/graphql")) + rest_request = authenticator(_prepared_request()) + authenticator.update_from_response( + graphql_request, _response(status_code=429, headers={"X-RateLimit-Remaining": "0"}) + ) + + assert authenticator._states["token_1"]["graphql"].remaining == 0 + assert authenticator.has_alternative_token(graphql_request) is True + # the rest pool is untouched, so nothing about it is "spent" and no rotation is implied + assert authenticator._states["token_1"]["rest"].remaining == 4999 + assert authenticator.has_alternative_token(rest_request) is False + + +def test_response_signed_by_something_else_is_not_attributed_to_a_token(requests_mock): + """The auth header may have been written by another component. Slicing off the configured + prefix without checking it would leave `_states` membership as the only guard against + charging the wrong token.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + authenticator._ensure_initialized() + request = _prepared_request() + request.headers["Authorization"] = "Bearer token_1" # different scheme + + authenticator.update_from_response(request, _response(headers={"X-RateLimit-Remaining": "1"})) + + assert authenticator._states["token_1"]["rest"].remaining == 5000 + + +def test_malformed_headers_are_ignored(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + request = authenticator(_prepared_request()) + + authenticator.update_from_response( + request, + _response(headers={"X-RateLimit-Remaining": "not-a-number", "X-RateLimit-Reset": "soon"}), + ) + + assert authenticator._states["token_1"]["rest"].remaining == 4999 + + +def test_concurrent_response_updates_do_not_lose_decrements(requests_mock): + """Clamp once, then charge a known number of calls with no further updates, so the assertion + fails if a single decrement is lost rather than only if the clamp itself is.""" + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _response_aware_authenticator(tokens=("token_1",)) + authenticator._ensure_initialized() + reset_at = authenticator._states["token_1"]["rest"].reset_at + seed_request = authenticator(_prepared_request()) + authenticator.update_from_response( + seed_request, + _response( + headers={ + "X-RateLimit-Remaining": "4000", + "X-RateLimit-Reset": str(int(reset_at.timestamp())), + } + ), + ) + assert authenticator._states["token_1"]["rest"].remaining == 4000 + + calls_per_thread, thread_count = 50, 4 + + def worker(): + for _ in range(calls_per_thread): + authenticator(_prepared_request()) + + threads = [threading.Thread(target=worker) for _ in range(thread_count)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert ( + authenticator._states["token_1"]["rest"].remaining == 4000 - calls_per_thread * thread_count + ) + + def _model(tokens): return RateLimitedMultipleTokenAuthenticatorModel.parse_obj( { @@ -422,3 +942,58 @@ def test_factory_interpolates_max_wait_time(): ) assert authenticator._max_wait_time == timedelta(minutes=30) + + +def test_factory_passes_response_header_fields_through(): + factory = ModelToComponentFactory() + model = _model("{{ config['pat'] }}") + model.quotas[0].remaining_header = "X-RateLimit-Remaining" + model.quotas[0].reset_header = "X-RateLimit-Reset" + model.quotas[0].limit_header = "X-RateLimit-Limit" + model.quotas[0].exhaustion_status_codes = [429] + + authenticator = factory.create_rate_limited_multiple_token_authenticator( + model, {"pat": "token_1"} + ) + + rest, graphql = authenticator._quotas + assert rest.remaining_header == "X-RateLimit-Remaining" + assert rest.reset_header == "X-RateLimit-Reset" + assert rest.limit_header == "X-RateLimit-Limit" + assert rest.exhaustion_status_codes == [429] + assert rest.is_response_aware + assert not graphql.is_response_aware + + +def test_factory_shares_instances_when_exhaustion_codes_are_omitted_versus_empty(): + """An omitted `exhaustion_status_codes` and an explicit `[]` behave identically at runtime, + so they must not key differently -- splitting the counters is the very failure this + component exists to avoid.""" + factory = ModelToComponentFactory() + config = {"pat": "token_1,token_2"} + + omitted = _model("{{ config['pat'] }}") + explicit_empty = _model("{{ config['pat'] }}") + explicit_empty.quotas[0].exhaustion_status_codes = [] + + assert factory.create_rate_limited_multiple_token_authenticator( + omitted, config + ) is factory.create_rate_limited_multiple_token_authenticator(explicit_empty, config) + + +def test_factory_does_not_share_instances_across_differing_header_config(): + """Header config is part of the authenticator's behaviour, so it must take part in the + instance cache key -- otherwise two differently configured pools would silently share state.""" + factory = ModelToComponentFactory() + config = {"pat": "token_1,token_2"} + + plain = factory.create_rate_limited_multiple_token_authenticator( + _model("{{ config['pat'] }}"), config + ) + response_aware_model = _model("{{ config['pat'] }}") + response_aware_model.quotas[0].remaining_header = "X-RateLimit-Remaining" + response_aware = factory.create_rate_limited_multiple_token_authenticator( + response_aware_model, config + ) + + assert plain is not response_aware diff --git a/unit_tests/sources/streams/http/test_http_client.py b/unit_tests/sources/streams/http/test_http_client.py index 48d396cb63..d7e0e51f81 100644 --- a/unit_tests/sources/streams/http/test_http_client.py +++ b/unit_tests/sources/streams/http/test_http_client.py @@ -2,6 +2,7 @@ import logging import os +import time from datetime import timedelta from unittest.mock import MagicMock, patch @@ -1062,6 +1063,261 @@ def update_response(*args, **kwargs): assert call_count == 2 +class _RecordingAuthenticator(TokenAuthenticator): + """An authenticator that tracks quota state and wants to see responses.""" + + def __init__(self): + super().__init__(token="token") + self.seen = [] + + def update_from_response(self, request, response): + self.seen.append(response.status_code) + + +def test_authenticator_receives_every_response(requests_mock): + """A quota-tracking authenticator only ever sees requests, so `HttpClient` hands it the + responses too. Every attempt counts, including the retried ones.""" + authenticator = _RecordingAuthenticator() + http_client = HttpClient( + name="test", + logger=logging.getLogger("test"), + authenticator=authenticator, + error_handler=HttpStatusErrorHandler(logger=logging.getLogger("test"), max_retries=1), + ) + requests_mock.get( + "https://example.com/", + [{"status_code": 500}, {"status_code": 200}], + ) + + with patch("time.sleep"): + http_client.send_request(http_method="GET", url="https://example.com/", request_kwargs={}) + + assert authenticator.seen == [500, 200] + + +def test_authenticator_without_update_hook_is_left_alone(requests_mock): + """Duck typed, so authenticators that don't track quota are untouched.""" + http_client = HttpClient( + name="test", + logger=logging.getLogger("test"), + authenticator=TokenAuthenticator(token="token"), + ) + requests_mock.get("https://example.com/", status_code=200) + + _, response = http_client.send_request( + http_method="GET", url="https://example.com/", request_kwargs={} + ) + + assert response.status_code == 200 + + +def test_authenticator_update_failure_does_not_break_the_request(requests_mock): + """Quota bookkeeping is best-effort; it must never turn a good response into a failure.""" + + class _BrokenAuthenticator(TokenAuthenticator): + def update_from_response(self, request, response): + raise ValueError("boom") + + http_client = HttpClient( + name="test", + logger=logging.getLogger("test"), + authenticator=_BrokenAuthenticator(token="token"), + ) + requests_mock.get("https://example.com/", status_code=200, json={"ok": True}) + + _, response = http_client.send_request( + http_method="GET", url="https://example.com/", request_kwargs={} + ) + + assert response.json() == {"ok": True} + + +def test_rate_limited_token_rotates_on_the_retry_end_to_end(requests_mock): + """End-to-end: a rate-limited response zeroes that token's pool and the retry goes out on + another token. + + This depends on `_send` re-signing the request on every attempt past the first. Nothing else + covers that line, so a refactor could silently drop rotation-on-retry and leave the sync + hammering a token the server has already rejected. + + The 30-minute reset in the response is deliberately far away: the assertion on `sleeps` is + that it is *not* waited out, because the authenticator reports another token with quota. + """ + from airbyte_cdk.sources.declarative.auth.rate_limited_multiple_token import ( + RateLimitedMultipleTokenAuthenticator, + TokenQuota, + ) + from airbyte_cdk.sources.declarative.requesters.error_handlers import DefaultErrorHandler + from airbyte_cdk.sources.declarative.requesters.error_handlers.backoff_strategies import ( + WaitUntilTimeFromHeaderBackoffStrategy, + ) + + quota_status_url = "https://api.example.com/rate_limit" + reset = int(time.time()) + 1800 + requests_mock.get( + quota_status_url, + json={"resources": {"core": {"remaining": 5000, "reset": reset, "limit": 5000}}}, + ) + authenticator = RateLimitedMultipleTokenAuthenticator( + tokens=["token_1", "token_2"], + quotas=[ + TokenQuota( + name="rest", + remaining_path=["resources", "core", "remaining"], + reset_path=["resources", "core", "reset"], + limit_path=["resources", "core", "limit"], + remaining_header="X-RateLimit-Remaining", + reset_header="X-RateLimit-Reset", + exhaustion_status_codes=[429], + ) + ], + quota_status_url=quota_status_url, + auth_method="token", + ) + + tokens_used = [] + + def respond(request, context): + token = request.headers["Authorization"].split()[-1] + tokens_used.append(token) + if token == "token_1": + context.status_code = 429 + context.headers = {"X-RateLimit-Reset": str(reset), "X-RateLimit-Remaining": "0"} + return "{}" + context.status_code = 200 + context.headers = {} + return "{}" + + requests_mock.get("https://api.example.com/data", text=respond) + # The backoff strategies must be on the client: that is what `_handle_error_resolution` + # consults. Passing them only to the error handler falls through to a 1s default and hides + # the real wait. + http_client = HttpClient( + name="test", + logger=logging.getLogger("test"), + authenticator=authenticator, + error_handler=DefaultErrorHandler(config={}, parameters={}), + backoff_strategy=[ + WaitUntilTimeFromHeaderBackoffStrategy( + header="X-RateLimit-Reset", config={}, parameters={} + ) + ], + ) + + sleeps = [] + with patch("time.sleep", side_effect=lambda seconds: sleeps.append(seconds)): + _, response = http_client.send_request( + http_method="GET", url="https://api.example.com/data", request_kwargs={} + ) + + assert response.status_code == 200 + assert tokens_used == ["token_1", "token_2"] + assert authenticator._states["token_1"]["rest"].remaining == 0 + # The 1800s reset is skipped: a token with quota was available the whole time. + assert max(sleeps) < 5, f"expected a prompt retry, slept {sleeps}" + + tokens_used.clear() + with patch("time.sleep", side_effect=lambda seconds: sleeps.append(seconds)): + http_client.send_request( + http_method="GET", url="https://api.example.com/data", request_kwargs={} + ) + assert tokens_used == ["token_2"] # the spent token is not tried again + + +class _SpareTokenAuthenticator(TokenAuthenticator): + """Reports a spare credential without tracking any quota.""" + + def __init__(self, has_spare=True): + super().__init__(token="token") + self.has_spare = has_spare + + def has_alternative_token(self, request): + return self.has_spare + + +def _rate_limited_client(authenticator, backoff_seconds=1800): + return HttpClient( + name="test", + logger=logging.getLogger("test"), + authenticator=authenticator, + error_handler=HttpStatusErrorHandler( + logger=logging.getLogger("test"), + error_mapping={ + 429: ErrorResolution( + response_action=ResponseAction.RATE_LIMITED, + failure_type=FailureType.transient_error, + error_message="rate limited", + ), + 500: ErrorResolution( + response_action=ResponseAction.RETRY, + failure_type=FailureType.transient_error, + error_message="server error", + ), + }, + max_retries=1, + ), + backoff_strategy=_ConstantBackoffStrategy(backoff_seconds), + ) + + +class _ConstantBackoffStrategy(BackoffStrategy): + def __init__(self, seconds): + self.seconds = seconds + + def backoff_time(self, *args, **kwargs): + return self.seconds + + +def test_rate_limit_wait_is_skipped_when_another_credential_is_available(requests_mock): + """The backoff is computed from the rejected credential's response, so it cannot know a + spare one is idle. The client asks before sleeping.""" + requests_mock.get("https://example.com/", [{"status_code": 429}, {"status_code": 200}]) + client = _rate_limited_client(_SpareTokenAuthenticator(has_spare=True)) + + sleeps = [] + with patch("time.sleep", side_effect=lambda seconds: sleeps.append(seconds)): + _, response = client.send_request( + http_method="GET", url="https://example.com/", request_kwargs={} + ) + + assert response.status_code == 200 + assert max(sleeps) < 5, f"expected a prompt retry, slept {sleeps}" + + +def test_rate_limit_wait_is_paid_when_no_other_credential_is_available(requests_mock): + requests_mock.get("https://example.com/", [{"status_code": 429}, {"status_code": 200}]) + client = _rate_limited_client(_SpareTokenAuthenticator(has_spare=False)) + + sleeps = [] + with patch("time.sleep", side_effect=lambda seconds: sleeps.append(seconds)): + client.send_request(http_method="GET", url="https://example.com/", request_kwargs={}) + + assert max(sleeps) > 1000, f"the full rate-limit wait should stand, slept {sleeps}" + + +def test_non_rate_limit_retry_is_not_shortened_by_a_spare_credential(requests_mock): + """A 500 has nothing to do with credentials; its backoff must be left alone.""" + requests_mock.get("https://example.com/", [{"status_code": 500}, {"status_code": 200}]) + client = _rate_limited_client(_SpareTokenAuthenticator(has_spare=True)) + + sleeps = [] + with patch("time.sleep", side_effect=lambda seconds: sleeps.append(seconds)): + client.send_request(http_method="GET", url="https://example.com/", request_kwargs={}) + + assert max(sleeps) > 1000, f"a server-error backoff should stand, slept {sleeps}" + + +def test_authenticator_without_alternative_token_hook_pays_the_full_wait(requests_mock): + requests_mock.get("https://example.com/", [{"status_code": 429}, {"status_code": 200}]) + client = _rate_limited_client(TokenAuthenticator(token="token")) + + sleeps = [] + with patch("time.sleep", side_effect=lambda seconds: sleeps.append(seconds)): + client.send_request(http_method="GET", url="https://example.com/", request_kwargs={}) + + assert max(sleeps) > 1000 + + def test_deprecated_alias_message_representation_airbyte_traced_errors_is_importable(): """Verify that the deprecated alias still resolves to AirbyteTracedException.""" assert MessageRepresentationAirbyteTracedErrors is AirbyteTracedException