From 07e1f7c830fbb38de7290932bfd8683656b1bc5b Mon Sep 17 00:00:00 2001 From: Daryna Ishchenko <80129833+darynaishchenko@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:23:24 +0300 Subject: [PATCH 1/2] feat(low-code): let a quota status endpoint report that rate limiting is off RateLimitedMultipleTokenAuthenticator seeds its counters from quota_status_url on the first signed request, through a per-token HttpClient built with no error handler -- so the default mapping applies and any non-2xx fails the connection before the connector issues a single stream request. That is right when the endpoint is expected to work, and wrong when it is optional. GitHub Enterprise Server leaves HTTP API rate limiting disabled by default, and an instance in that state answers GET /rate_limit with 404 "Rate limiting is not enabled." (see vermiculus/magithub#104, and actions/stale#1227 for the same message on GHES 3.15). Every command then dies at seeding with a generic "Resource not found", and the connector has no workaround: QuotaStatusSource exposes only url, http_method and request_headers, and a stream's error_handler does not apply to the authenticator's own client. QuotaStatusSource gains an opt-in `unavailable_status_codes`. Listed statuses map to ResponseAction.IGNORE for the quota request only, so send_request returns the response instead of raising, and every pool is seeded untracked. An untracked pool skips exhaustion waits, proactive throttling and rotation-on-exhaustion, while requests are still signed. A missing quota path in an otherwise healthy response is treated the same way, per pool, rather than failing the connection -- also only under the opt-in, so connectors that expect the path keep the loud config error. A list rather than a hardcoded 404 because GitHub documents 404 as a possible response for this endpoint but nowhere states that it means rate limiting is disabled, and a proxy in front of an instance can answer differently. The undocumented assumption belongs in the connector that makes it. Untracked is a `tracked: bool` on _QuotaState rather than a very large `remaining`: six call sites read that state and a sentinel would have to satisfy all of them by arithmetic accident, and the far-future reset_at it would need makes both branches of update_from_response unreachable, silently discarding response headers on deployments that do send them. Deliberately narrow: this suppresses only the authenticator's own bookkeeping. GHES exposes HTTP API and secondary rate limiting as independent toggles, so an instance can 404 the quota endpoint and still answer 403 when pushed -- responses that report a rate limit remain the stream error handler's job, and the tests pin that separation. No behaviour change for any connector that does not set the field. Co-Authored-By: Claude Opus 5 (1M context) --- .../auth/rate_limited_multiple_token.py | 126 +++++++++++++- .../declarative_component_schema.yaml | 8 + .../models/declarative_component_schema.py | 6 + .../parsers/model_to_component_factory.py | 7 + .../auth/test_rate_limited_multiple_token.py | 164 ++++++++++++++++++ 5 files changed, 308 insertions(+), 3 deletions(-) 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 f56797109..55a8d7407 100644 --- a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py +++ b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py @@ -16,6 +16,14 @@ from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator from airbyte_cdk.sources.streams.call_rate import RequestMatcher from airbyte_cdk.sources.streams.http import HttpClient +from airbyte_cdk.sources.streams.http.error_handlers import HttpStatusErrorHandler +from airbyte_cdk.sources.streams.http.error_handlers.default_error_mapping import ( + DEFAULT_ERROR_MAPPING, +) +from airbyte_cdk.sources.streams.http.error_handlers.response_models import ( + ErrorResolution, + ResponseAction, +) from airbyte_cdk.sources.streams.http.requests_native_auth import TokenAuthenticator from airbyte_cdk.utils import AirbyteTracedException from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now, ab_datetime_parse @@ -56,11 +64,33 @@ def is_response_aware(self) -> bool: ) +class _Missing: + """Marks a quota path the response did not contain. Distinct from `None`, which a server is + free to send as a legitimate value.""" + + def __repr__(self) -> str: + return "" + + +_MISSING = _Missing() + + @dataclass class _QuotaState: remaining: int reset_at: AirbyteDateTime limit: int + tracked: bool = True + """Whether the server reports a quota for this pool at all. + + False means the quota status endpoint answered with one of `unavailable_status_codes`, or + omitted this pool's path from an otherwise-healthy response. The pool then has no numbers + worth acting on, so every decision derived from them is skipped rather than made against + invented ones. Modelled as a flag rather than a very large `remaining` because six call + sites read this state and a sentinel would have to satisfy all of them by arithmetic + accident -- and a far-future `reset_at` would silently make both branches of + `update_from_response` unreachable, discarding real headers if the server sends them. + """ class RateLimitedMultipleTokenAuthenticator(DeclarativeAuthenticator): @@ -104,6 +134,7 @@ def __init__( quota_status_url: str, quota_status_http_method: str = "GET", quota_status_headers: Optional[Mapping[str, str]] = None, + quota_status_unavailable_status_codes: Optional[List[int]] = None, auth_method: str = "Bearer", header: str = "Authorization", max_wait_time: timedelta = timedelta(hours=2), @@ -134,11 +165,14 @@ def __init__( self._budget_reserve_fraction = budget_reserve_fraction self._budget_min_reserve = budget_min_reserve + self._unavailable_status_codes = set(quota_status_unavailable_status_codes or []) + self._lock = threading.RLock() self._refresh_lock = threading.Lock() self._initialized = False self._budget_logged = False self._unmatched_logged = False + self._untracked_logged = False self._states: dict[str, dict[str, _QuotaState]] = {} self._token_to_http_client: Mapping[str, HttpClient] = { token: HttpClient( @@ -148,6 +182,7 @@ def __init__( token, auth_method=self._auth_method, auth_header=self._header ), use_cache=False, # quota values change frequently; never reuse cached responses + error_handler=self._quota_status_error_handler(), ) for token in self._tokens } @@ -171,6 +206,54 @@ def __call__(self, request: requests.PreparedRequest) -> Any: request.headers[self._header] = f"{self._auth_method} {token}".strip() return request + def _quota_status_error_handler(self) -> Optional[HttpStatusErrorHandler]: + """Error handling for the quota status request itself. + + `None` keeps `HttpClient`'s default, under which every non-2xx fails the connection -- + which is correct when the endpoint is expected to work. When the connector has declared + that some statuses mean "quota tracking is not enabled here", those are mapped to + `IGNORE` instead, so `send_request` hands the response back rather than raising and + `_fetch_quota_states` can decide what it means. Statuses outside the list keep failing. + """ + if not self._unavailable_status_codes: + return None + return HttpStatusErrorHandler( + self._logger, + error_mapping={ + **DEFAULT_ERROR_MAPPING, + **{ + status_code: ErrorResolution( + response_action=ResponseAction.IGNORE, + failure_type=FailureType.transient_error, + error_message=( + "Quota status endpoint reports that rate limiting is unavailable; " + "treating token quotas as untracked." + ), + ) + for status_code in self._unavailable_status_codes + }, + }, + ) + + def _untracked_states(self) -> dict[str, _QuotaState]: + """A state per pool meaning "the server tracks nothing here".""" + now = ab_datetime_now() + return { + quota.name: _QuotaState(remaining=0, reset_at=now, limit=0, tracked=False) + for quota in self._quotas + } + + def _log_untracked_once(self, reason: str) -> None: + if self._untracked_logged: + return + self._untracked_logged = True + self._logger.info( + "Quota status endpoint %s. Token quotas are untracked: the connector will not wait " + "for quota resets, throttle proactively, or rotate tokens on exhaustion. Responses " + "that report a rate limit are still handled by the stream's error handler.", + reason, + ) + def _ensure_initialized(self) -> None: if self._initialized: return @@ -210,10 +293,19 @@ def _acquire_call(self, quota: TokenQuota) -> str: with self._lock: token = self._active_token state = self._states[token][quota.name] + if not state.tracked: + # Nothing to spend and nothing to wait for. Rotation still works on demand + # (`update_from_response`, an explicit exhaustion signal), it just is not + # driven by counters that do not exist. + return token if state.remaining > 0: state.remaining -= 1 budget_delay = self._compute_budget_delay(quota) - elif all(self._states[token][quota.name].remaining <= 0 for token in self._tokens): + elif all( + self._states[token][quota.name].remaining <= 0 + and self._states[token][quota.name].tracked + for token in self._tokens + ): now = time.monotonic() if exhaustion_deadline is None: exhaustion_deadline = now + self._max_wait_time.total_seconds() @@ -262,6 +354,8 @@ def _acquire_call(self, quota: TokenQuota) -> str: def _compute_budget_delay(self, quota: TokenQuota) -> Optional[float]: """Compute the proactive throttling delay. Must be called while holding the lock.""" states = [self._states[token][quota.name] for token in self._tokens] + if any(not state.tracked for state in states): + return None if not all(state.remaining <= self._get_budget_reserve(state) for state in states): return None @@ -295,7 +389,9 @@ def _refresh_after_exhaustion(self, quota: TokenQuota) -> None: with self._refresh_lock: with self._lock: still_exhausted = all( - self._states[token][quota.name].remaining <= 0 for token in self._tokens + self._states[token][quota.name].remaining <= 0 + and self._states[token][quota.name].tracked + for token in self._tokens ) if still_exhausted: self._seed_all_tokens() @@ -320,6 +416,11 @@ def _fetch_quota_states(self, token: str) -> dict[str, _QuotaState]: headers=self._quota_status_headers, request_kwargs={}, ) + if response.status_code in self._unavailable_status_codes: + # Only reachable when the connector opted in: without `unavailable_status_codes` + # the default error mapping raises before this point. + self._log_untracked_once(f"returned HTTP {response.status_code}") + return self._untracked_states() response_body = response.json() states = {} @@ -331,6 +432,15 @@ def _fetch_quota_states(self, token: str) -> dict[str, _QuotaState]: if quota.limit_path else remaining ) + if remaining is _MISSING or reset is _MISSING or limit is _MISSING: + # A deployment that reports some pools but not others. Losing the whole + # connection over one absent key is worse than running that pool untracked -- + # but only for connectors that opted into tolerating this endpoint at all. + self._log_untracked_once( + f"did not report quota '{quota.name}'; that pool is untracked" + ) + states[quota.name] = self._untracked_states()[quota.name] + continue states[quota.name] = _QuotaState( remaining=int(remaining), reset_at=ab_datetime_parse(reset), @@ -342,6 +452,8 @@ def _extract_path(self, response_body: Mapping[str, Any], path: List[str]) -> An value: Any = response_body for key in path: if not isinstance(value, Mapping) or key not in value: + if self._unavailable_status_codes: + return _MISSING raise AirbyteTracedException( failure_type=FailureType.config_error, internal_message=f"Quota status response did not contain expected path: {path}", @@ -383,6 +495,12 @@ def update_from_response( 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 not state.tracked: + # The quota status endpoint said this pool is not tracked. Response headers + # could contradict that, but adopting them would resurrect exhaustion waits and + # throttling on a deployment that has rate limiting switched off. Rate-limit + # *responses* remain the error handler's job either way. + return 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 @@ -438,10 +556,12 @@ def has_alternative_token(self, request: requests.PreparedRequest) -> bool: with self._lock: if not self._states or sender is None: return False - if self._states[sender][quota.name].remaining > 0: + sender_state = self._states[sender][quota.name] + if not sender_state.tracked or sender_state.remaining > 0: return False return any( self._states[token][quota.name].remaining > 0 + and self._states[token][quota.name].tracked for token in self._tokens if token != sender ) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 091c694d2..aa481d58c 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -444,6 +444,14 @@ definitions: type: object additionalProperties: type: string + unavailable_status_codes: + title: Unavailable Status Codes + description: Status codes from the quota status endpoint that mean quota tracking is unavailable rather than broken, such as a self-hosted deployment with rate limiting turned off. Every pool is then treated as untracked, so the authenticator stops waiting for quota resets, stops throttling proactively and stops rotating on exhaustion, while still signing requests. Rate limiting reported by ordinary responses is unaffected, so an error handler that retries 429 or 403 keeps working. Any status not listed still fails the connection, so list only the codes the endpoint uses to report that rate limiting is not enabled. + type: array + items: + type: integer + examples: + - [404] $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 37eb50b7b..e34f2c4c1 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -575,6 +575,12 @@ class QuotaStatusSource(BaseModel): description="Additional headers to send with the quota status request.", title="Request Headers", ) + unavailable_status_codes: Optional[List[int]] = Field( + None, + description="Status codes from the quota status endpoint that mean quota tracking is unavailable rather than broken, such as a self-hosted deployment with rate limiting turned off. Every pool is then treated as untracked, so the authenticator stops waiting for quota resets, stops throttling proactively and stops rotating on exhaustion, while still signing requests. Rate limiting reported by ordinary responses is unaffected, so an error handler that retries 429 or 403 keeps working. Any status not listed still fails the connection, so list only the codes the endpoint uses to report that rate limiting is not enabled.", + examples=[[404]], + title="Unavailable 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 a96c406d8..b36d4661e 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4693,6 +4693,11 @@ def create_rate_limited_multiple_token_authenticator( key: str(InterpolatedString.create(value, parameters={}).eval(config)) for key, value in (model.quota_status_source.request_headers or {}).items() } + # Normalized the same way as the quota specs above, so an omitted field and an explicit + # `[]` key identically and keep sharing one set of counters. + quota_status_unavailable_status_codes = sorted( + model.quota_status_source.unavailable_status_codes or [] + ) auth_method = model.auth_method or "Bearer" header = model.header or "Authorization" max_wait_time_str = str( @@ -4723,6 +4728,7 @@ def create_rate_limited_multiple_token_authenticator( "quota_status_url": quota_status_url, "quota_status_http_method": quota_status_http_method, "quota_status_headers": quota_status_headers, + "quota_status_unavailable_status_codes": quota_status_unavailable_status_codes, "auth_method": auth_method, "header": header, "max_wait_time": max_wait_time.total_seconds(), @@ -4758,6 +4764,7 @@ def create_rate_limited_multiple_token_authenticator( quota_status_url=quota_status_url, quota_status_http_method=quota_status_http_method, quota_status_headers=quota_status_headers, + quota_status_unavailable_status_codes=quota_status_unavailable_status_codes, auth_method=auth_method, header=header, max_wait_time=max_wait_time, 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 0a4580091..461476746 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 @@ -997,3 +997,167 @@ def test_factory_does_not_share_instances_across_differing_header_config(): ) assert plain is not response_aware + + +def test_unavailable_status_is_untracked_and_never_blocks(requests_mock): + """A deployment with rate limiting switched off answers the quota endpoint with an error. + + Opting in must turn that into "there is no quota here" rather than a failed connection: + requests are still signed, nothing waits for a reset that will never come, and the + proactive budget never throttles. + """ + requests_mock.get( + QUOTA_STATUS_URL, status_code=404, json={"message": "Rate limiting is not enabled."} + ) + authenticator = _authenticator(quota_status_unavailable_status_codes=[404]) + + with patch("time.sleep") as sleep_mock: + for _ in range(3): + request = authenticator(_prepared_request()) + + assert request.headers["Authorization"] == "token token_1" + sleep_mock.assert_not_called() + for token in ("token_1", "token_2"): + for pool in ("rest", "graphql"): + assert authenticator._states[token][pool].tracked is False + + +def test_unavailable_status_without_opt_in_still_fails(requests_mock): + """Unchanged behaviour for every connector that has not opted in -- the endpoint failing is + still a broken connection, not a silent switch to untracked quotas.""" + requests_mock.get( + QUOTA_STATUS_URL, status_code=404, json={"message": "Rate limiting is not enabled."} + ) + authenticator = _authenticator() + + with pytest.raises(AirbyteTracedException): + authenticator(_prepared_request()) + + +def test_status_outside_the_opt_in_list_still_fails(requests_mock): + """The opt-in is a list of specific codes, not blanket tolerance -- a 500 from the quota + endpoint is a real failure even when 404 is excused.""" + requests_mock.get(QUOTA_STATUS_URL, status_code=500, json={"message": "Internal Server Error"}) + authenticator = _authenticator(quota_status_unavailable_status_codes=[404]) + + with pytest.raises(AirbyteTracedException): + authenticator(_prepared_request()) + + +def test_untracked_pool_reports_no_alternative_token(requests_mock): + """`has_alternative_token` answers "should HttpClient skip the rate-limit wait and retry on + another credential". With no counters it cannot claim a token is spent, so it must say no + and let the computed backoff stand.""" + requests_mock.get( + QUOTA_STATUS_URL, status_code=404, json={"message": "Rate limiting is not enabled."} + ) + authenticator = _authenticator(quota_status_unavailable_status_codes=[404]) + request = _prepared_request() + authenticator(request) + + assert authenticator.has_alternative_token(request) is False + + +def test_untracked_pool_ignores_response_headers(requests_mock): + """A pool the quota endpoint does not track stays untracked even if responses carry quota + headers. Adopting them would resurrect exhaustion waits and throttling on a deployment that + deliberately has rate limiting turned off; responses that actually report a rate limit are + the error handler's job.""" + requests_mock.get( + QUOTA_STATUS_URL, status_code=404, json={"message": "Rate limiting is not enabled."} + ) + quotas = [ + TokenQuota( + name="rest", + remaining_path=["resources", "core", "remaining"], + reset_path=["resources", "core", "reset"], + remaining_header="X-RateLimit-Remaining", + reset_header="X-RateLimit-Reset", + ) + ] + authenticator = RateLimitedMultipleTokenAuthenticator( + tokens=["token_1"], + quotas=quotas, + quota_status_url=QUOTA_STATUS_URL, + quota_status_unavailable_status_codes=[404], + auth_method="token", + ) + request = _prepared_request() + authenticator(request) + + response = requests.Response() + response.status_code = 200 + response.headers["X-RateLimit-Remaining"] = "17" + response.headers["X-RateLimit-Reset"] = str(int(time.time()) + 3600) + authenticator.update_from_response(request, response) + + assert authenticator._states["token_1"]["rest"].tracked is False + + +def test_missing_quota_path_is_untracked_only_for_that_pool(requests_mock): + """A deployment that reports some pools but not others should lose the pool, not the + connection -- and only when the connector has opted into tolerating this endpoint.""" + body = _quota_status_body() + del body["resources"]["graphql"] + requests_mock.get(QUOTA_STATUS_URL, json=body) + authenticator = _authenticator(tokens=("token_1",), quota_status_unavailable_status_codes=[404]) + + authenticator(_prepared_request()) + + assert authenticator._states["token_1"]["graphql"].tracked is False + assert authenticator._states["token_1"]["rest"].tracked is True + assert authenticator._states["token_1"]["rest"].remaining == 4999 + + +def test_missing_quota_path_without_opt_in_still_raises(requests_mock): + body = _quota_status_body() + del body["resources"]["graphql"] + requests_mock.get(QUOTA_STATUS_URL, json=body) + authenticator = _authenticator(tokens=("token_1",)) + + with pytest.raises(AirbyteTracedException): + authenticator(_prepared_request()) + + +def test_unavailable_status_codes_are_threaded_through_the_factory(): + """The manifest field has to reach the constructor, and two definitions that differ only by + it must not collide in the factory's instance cache.""" + definition = { + "type": "RateLimitedMultipleTokenAuthenticator", + "tokens": "token_1,token_2", + "token_delimiter": ",", + "quota_status_source": { + "type": "QuotaStatusSource", + "url": QUOTA_STATUS_URL, + "unavailable_status_codes": [404], + }, + "quotas": [ + { + "type": "TokenQuota", + "name": "rest", + "remaining_path": ["resources", "core", "remaining"], + "reset_path": ["resources", "core", "reset"], + } + ], + } + factory = ModelToComponentFactory() + transformer = ManifestComponentTransformer() + + def build(component_definition): + propagated = transformer.propagate_types_and_parameters("", component_definition, {}) + return factory.create_component( + model_type=RateLimitedMultipleTokenAuthenticatorModel, + component_definition=propagated, + config={}, + ) + + tolerant = build(definition) + assert tolerant._unavailable_status_codes == {404} + + without = { + **definition, + "quota_status_source": {"type": "QuotaStatusSource", "url": QUOTA_STATUS_URL}, + } + strict = build(without) + assert strict._unavailable_status_codes == set() + assert strict is not tolerant, "the cache key must include the new field" From d501e4f606d5b993d88f4d9feb1ef40b47cf71e3 Mon Sep 17 00:00:00 2001 From: Daryna Ishchenko <80129833+darynaishchenko@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:17:27 +0300 Subject: [PATCH 2/2] feat(low-code): narrow unavailable_status_codes and spread load across untracked tokens Addresses review on 07e1f7c8. unavailable_status_codes now gates status codes only. A quota path the response does not contain raises again, opt-in or not: an endpoint answering with an error is telling you it does not track quotas, while an endpoint answering with a body does track them, so a path absent from that body is a wrong path. Letting one field cover both meant a typo in remaining_path silently switched quota tracking off for the whole sync. The _Missing sentinel and its per-pool untrack branch are gone with it, which also removes the limit_path asymmetry (a declared-but-absent limit_path untracked a pool whose remaining and reset were both present) and the per-pool reseed that could flip an untracked pool back to tracked mid-sync. That failure is reclassified config_error -> system_error. The quota paths come from the manifest, so there is nothing in the user's configuration to correct. The message now names the pool and which of remaining/reset/limit was absent. Untracked tokens are now rotated round-robin. Every token hits the same quota_status_url and gets the same status, so on a deployment that reports no quota the untracked branch is the only one _acquire_call ever takes; without advancing the active token there, one credential served the entire sync and the rest of a multi-token configuration went unused. Nothing about the quota endpoint being unavailable implies the other credentials should sit idle, and the server may still enforce limits it declines to report. has_alternative_token still answers False for an untracked sender, but the docstring's old reason no longer held once the retry started rotating. The real reason is narrower: what it withholds is the skipped wait, the backoff it would skip is computed from the server's own reset header, and an untracked pool has no counters with which to argue the rejection was about that credential. Overriding that on a guess would burn every retry in under a second whenever the limit turns out to be shared across credentials. The untracked summary log moved to _seed_all_tokens, the first point that can see every token, so it states the scope of the consequence instead of asserting a global one while another token is still tracked and still throttling. Its partial-case wording says what the tracked tokens actually keep: they throttle until their counters are locally spent and are then not refreshed, because the exhaustion wait is the only reseed after startup and it is unreachable while any token is untracked. The custom error_message on the IGNORE resolution is dropped so HttpClient logs its own per-request line rather than a near-duplicate of the summary. The factory deduplicates as well as sorts the status codes, and the schema field carries uniqueItems: true, following HttpResponseFilter.http_codes. Without it [404] and [404, 404] built two authenticators that did not share quota counters while behaving identically. The field description also warns against listing authentication or authorization statuses, which would read a revoked credential as quota tracking being unavailable. Tests: the untracked-headers test asserted tracked is False, which update_from_response never writes, so it passed with its guard deleted; it now asserts the counters the guard protects. Four tests cover the mixed tracked/untracked state, the only state in which four of the tracked guards are reachable. Added: round-robin across untracked tokens, an untracked pool is never reseeded, duplicate status codes are rejected, and order-only differences share one set of counters. Each of the seven guards was checked by deleting it and confirming the suite fails. 166 passed in the authenticator suite, 795 across declarative/auth, declarative/parsers and streams/http. ruff, format and mypy clean. Co-Authored-By: Claude Opus 5 (1M context) --- .../auth/rate_limited_multiple_token.py | 172 ++++++++---- .../declarative_component_schema.yaml | 3 +- .../models/declarative_component_schema.py | 3 +- .../parsers/model_to_component_factory.py | 10 +- .../auth/test_rate_limited_multiple_token.py | 260 ++++++++++++++++-- 5 files changed, 366 insertions(+), 82 deletions(-) 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 55a8d7407..064e7d6b7 100644 --- a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py +++ b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py @@ -64,17 +64,6 @@ def is_response_aware(self) -> bool: ) -class _Missing: - """Marks a quota path the response did not contain. Distinct from `None`, which a server is - free to send as a legitimate value.""" - - def __repr__(self) -> str: - return "" - - -_MISSING = _Missing() - - @dataclass class _QuotaState: remaining: int @@ -83,13 +72,16 @@ class _QuotaState: tracked: bool = True """Whether the server reports a quota for this pool at all. - False means the quota status endpoint answered with one of `unavailable_status_codes`, or - omitted this pool's path from an otherwise-healthy response. The pool then has no numbers - worth acting on, so every decision derived from them is skipped rather than made against - invented ones. Modelled as a flag rather than a very large `remaining` because six call - sites read this state and a sentinel would have to satisfy all of them by arithmetic - accident -- and a far-future `reset_at` would silently make both branches of - `update_from_response` unreachable, discarding real headers if the server sends them. + False means the quota status endpoint answered with one of `unavailable_status_codes`. The + pool then has no numbers worth acting on, so every decision derived from them is skipped + rather than made against invented ones. A quota *path* missing from an otherwise-healthy + response is not this case and still fails: the endpoint answering at all means it reports + quotas, so an absent path is a wrong path. + + Modelled as a flag rather than a very large `remaining` because six call sites read this + state and a sentinel would have to satisfy all of them by arithmetic accident -- and a + far-future `reset_at` would silently make both branches of `update_from_response` + unreachable, discarding real headers if the server sends them. """ @@ -225,10 +217,6 @@ def _quota_status_error_handler(self) -> Optional[HttpStatusErrorHandler]: status_code: ErrorResolution( response_action=ResponseAction.IGNORE, failure_type=FailureType.transient_error, - error_message=( - "Quota status endpoint reports that rate limiting is unavailable; " - "treating token quotas as untracked." - ), ) for status_code in self._unavailable_status_codes }, @@ -236,23 +224,62 @@ def _quota_status_error_handler(self) -> Optional[HttpStatusErrorHandler]: ) def _untracked_states(self) -> dict[str, _QuotaState]: - """A state per pool meaning "the server tracks nothing here".""" + """A state per pool meaning "the server tracks nothing here". + + `remaining=0` is load-bearing rather than arbitrary: it is what keeps every + `remaining > 0` test in this class correct for an untracked pool without also having to + consult `tracked`. Nothing ever raises it, since `_acquire_call` only decrements and + `update_from_response` returns early for an untracked pool. + """ now = ab_datetime_now() return { quota.name: _QuotaState(remaining=0, reset_at=now, limit=0, tracked=False) for quota in self._quotas } - def _log_untracked_once(self, reason: str) -> None: + def _log_untracked_tokens(self, states: Mapping[str, Mapping[str, _QuotaState]]) -> None: + """Report untracked tokens once, scoped to how many of them there are. + + Deliberately called with every token's states rather than from `_fetch_quota_states`, + which sees one token at a time. The consequence of untracking -- no exhaustion waits, no + proactive throttling, no rotation -- is only true of the tokens that are untracked, and + a per-token call site cannot know whether the others are. Claiming it globally while one + token is still tracked and still doing all three would send an operator looking for a + problem in the wrong place. + """ if self._untracked_logged: return + untracked = [ + token + for token, pools in states.items() + if any(not state.tracked for state in pools.values()) + ] + if not untracked: + return self._untracked_logged = True - self._logger.info( - "Quota status endpoint %s. Token quotas are untracked: the connector will not wait " - "for quota resets, throttle proactively, or rotate tokens on exhaustion. Responses " - "that report a rate limit are still handled by the stream's error handler.", - reason, - ) + if len(untracked) == len(self._tokens): + self._logger.info( + "Quota status endpoint reports that rate limiting is unavailable. Token quotas " + "are untracked: the connector will not wait for quota resets, throttle " + "proactively, or rotate tokens on exhaustion. Responses that report a rate " + "limit are still handled by the stream's error handler." + ) + else: + # Not "the others are unaffected": `_acquire_call` rotates onto an untracked token + # rather than waiting, so the exhaustion wait -- and with it the only reseed after + # startup -- becomes unreachable as soon as one token is untracked. The tracked + # tokens keep throttling until their counters are locally spent and are then left + # spent for the rest of the sync. + self._logger.info( + "Quota status endpoint reports that rate limiting is unavailable for %d of %d " + "tokens. Those tokens are used without quota tracking. The other %d keep " + "proactive throttling until their counters are locally spent, after which " + "traffic moves onto the untracked tokens: the connector no longer waits for a " + "quota reset, so it never refreshes them.", + len(untracked), + len(states), + len(states) - len(untracked), + ) def _ensure_initialized(self) -> None: if self._initialized: @@ -294,9 +321,20 @@ def _acquire_call(self, quota: TokenQuota) -> str: token = self._active_token state = self._states[token][quota.name] if not state.tracked: - # Nothing to spend and nothing to wait for. Rotation still works on demand - # (`update_from_response`, an explicit exhaustion signal), it just is not - # driven by counters that do not exist. + # Nothing to spend and nothing to wait for, but the tokens are still there + # to spread load over. Every token hits the same `quota_status_url` and so + # gets the same status, which means this branch is the *only* one taken on a + # deployment that reports no quota -- so without advancing here, one + # credential would serve the entire sync and the rest would go unused. + # Round-robin is the right rule precisely because there are no counters: + # nothing distinguishes the tokens, and the server may still enforce limits + # it declines to report. + # + # Note the other half of the mechanism: once any token is untracked the + # exhaustion branch below can never fire, so `_refresh_after_exhaustion` -- + # the only reseed after startup -- is unreachable, and a tracked token's + # quota is never picked up again even after its window resets. + self._active_token = next(self._tokens_iter) return token if state.remaining > 0: state.remaining -= 1 @@ -385,7 +423,13 @@ def _sleep_with_heartbeat(self, total_seconds: float, quota_name: str) -> None: ) def _refresh_after_exhaustion(self, quota: TokenQuota) -> None: - """Refresh counters after an exhaustion wait. Only one thread refreshes; others re-check state.""" + """Refresh counters after an exhaustion wait. Only one thread refreshes; others re-check state. + + The `tracked` term is not reachable from a single-threaded run -- reaching the wait at + all requires every token to be tracked -- but it is reachable under concurrency, because + another thread's reseed can untrack a token while this one sleeps. Reseeding then buys + nothing: `_acquire_call` will rotate onto the untracked token instead of waiting again. + """ with self._refresh_lock: with self._lock: still_exhausted = all( @@ -407,6 +451,7 @@ def _seed_all_tokens(self) -> None: with self._lock: self._states = states self._budget_logged = False + self._log_untracked_tokens(states) def _fetch_quota_states(self, token: str) -> dict[str, _QuotaState]: http_client = self._token_to_http_client[token] @@ -418,29 +463,23 @@ def _fetch_quota_states(self, token: str) -> dict[str, _QuotaState]: ) if response.status_code in self._unavailable_status_codes: # Only reachable when the connector opted in: without `unavailable_status_codes` - # the default error mapping raises before this point. - self._log_untracked_once(f"returned HTTP {response.status_code}") + # the default error mapping raises before this point. `_seed_all_tokens` reports it + # once every token has been fetched, which is the first point at which the scope of + # the consequence is known. return self._untracked_states() response_body = response.json() states = {} for quota in self._quotas: - remaining = self._extract_path(response_body, quota.remaining_path) - reset = self._extract_path(response_body, quota.reset_path) + remaining = self._extract_path( + response_body, quota.remaining_path, quota.name, "remaining" + ) + reset = self._extract_path(response_body, quota.reset_path, quota.name, "reset") limit = ( - self._extract_path(response_body, quota.limit_path) + self._extract_path(response_body, quota.limit_path, quota.name, "limit") if quota.limit_path else remaining ) - if remaining is _MISSING or reset is _MISSING or limit is _MISSING: - # A deployment that reports some pools but not others. Losing the whole - # connection over one absent key is worse than running that pool untracked -- - # but only for connectors that opted into tolerating this endpoint at all. - self._log_untracked_once( - f"did not report quota '{quota.name}'; that pool is untracked" - ) - states[quota.name] = self._untracked_states()[quota.name] - continue states[quota.name] = _QuotaState( remaining=int(remaining), reset_at=ab_datetime_parse(reset), @@ -448,16 +487,30 @@ def _fetch_quota_states(self, token: str) -> dict[str, _QuotaState]: ) return states - def _extract_path(self, response_body: Mapping[str, Any], path: List[str]) -> Any: + def _extract_path( + self, response_body: Mapping[str, Any], path: List[str], quota_name: str, field_name: str + ) -> Any: + """Read a configured quota path out of the response, or fail. + + A path the response does not contain is a `system_error` rather than a `config_error`: + the paths come from the manifest, not from anything the end user can edit, so there is + no configuration for them to correct. `unavailable_status_codes` does not soften this -- + it says what an endpoint answering with an error *means*, and an endpoint that answers + with a body does report quotas, so a path missing from that body is a wrong path. + """ value: Any = response_body for key in path: if not isinstance(value, Mapping) or key not in value: - if self._unavailable_status_codes: - return _MISSING raise AirbyteTracedException( - failure_type=FailureType.config_error, - internal_message=f"Quota status response did not contain expected path: {path}", - message="Quota status response is missing an expected field.", + failure_type=FailureType.system_error, + internal_message=( + f"Quota status response did not contain the {field_name} path {path} " + f"configured for quota '{quota_name}'" + ), + message=( + f"Quota status response does not contain the configured {field_name} " + f'path for token quota "{quota_name}".' + ), ) value = value[key] return value @@ -550,6 +603,16 @@ def has_alternative_token(self, request: requests.PreparedRequest) -> bool: 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. + + An untracked sender answers False too, but for a different reason, and it is a trade-off + rather than a clear win. The retry does rotate -- `_acquire_call` round-robins untracked + tokens -- so what this withholds is only the *skipped wait*. The backoff it would skip is + computed from what the server said (a reset or `Retry-After` header), and an untracked + pool has no counters with which to argue the rejection was about this credential + specifically. Overriding the server's own instruction on a guess would, when the limit is + shared across credentials, burn every retry in under a second and fail a request that + waiting would have completed. So a rate-limited response on an untracked pool rotates + credentials but still pays the computed backoff. """ quota = self._match_quota(request) sender = self._token_from_request(request) @@ -561,7 +624,6 @@ def has_alternative_token(self, request: requests.PreparedRequest) -> bool: return False return any( self._states[token][quota.name].remaining > 0 - and self._states[token][quota.name].tracked for token in self._tokens if token != sender ) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index aa481d58c..e3e36ac56 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -446,10 +446,11 @@ definitions: type: string unavailable_status_codes: title: Unavailable Status Codes - description: Status codes from the quota status endpoint that mean quota tracking is unavailable rather than broken, such as a self-hosted deployment with rate limiting turned off. Every pool is then treated as untracked, so the authenticator stops waiting for quota resets, stops throttling proactively and stops rotating on exhaustion, while still signing requests. Rate limiting reported by ordinary responses is unaffected, so an error handler that retries 429 or 403 keeps working. Any status not listed still fails the connection, so list only the codes the endpoint uses to report that rate limiting is not enabled. + description: Status codes from the quota status endpoint that mean quota tracking is unavailable rather than broken, such as a self-hosted deployment with rate limiting turned off. Every pool of the token whose request returned that status is then treated as untracked, so the authenticator stops waiting for quota resets, stops throttling proactively and stops rotating on exhaustion for it, while still signing requests. A token untracked this way stays untracked for the rest of the sync, because the endpoint is never consulted for it again, so a status the endpoint can also return transiently costs quota tracking for the whole run. If only some tokens return that status the others stay tracked, but they are no longer refreshed either, because the authenticator stops waiting for quota resets as soon as one token is untracked; once their counters are locally spent all traffic moves onto the untracked tokens. Rate limiting reported by ordinary responses is still handled by the stream's error handler, so one that retries 429 or 403 keeps working, and a retry rotates onto the next token; it pays the backoff the response asks for rather than the shortened one a tracked pool would get, since an untracked pool has no counters with which to argue the rejection was about that credential. Any status not listed still fails the connection, and this field never excuses a quota path missing from a response the endpoint did answer, so list only the codes the endpoint uses to report that rate limiting is not enabled. Do not list authentication or authorization statuses, since a 401 or 403 from a revoked credential would then be read as quota tracking being unavailable rather than as a credentials failure. type: array items: type: integer + uniqueItems: true examples: - [404] $parameters: diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index e34f2c4c1..944375e22 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -577,9 +577,10 @@ class QuotaStatusSource(BaseModel): ) unavailable_status_codes: Optional[List[int]] = Field( None, - description="Status codes from the quota status endpoint that mean quota tracking is unavailable rather than broken, such as a self-hosted deployment with rate limiting turned off. Every pool is then treated as untracked, so the authenticator stops waiting for quota resets, stops throttling proactively and stops rotating on exhaustion, while still signing requests. Rate limiting reported by ordinary responses is unaffected, so an error handler that retries 429 or 403 keeps working. Any status not listed still fails the connection, so list only the codes the endpoint uses to report that rate limiting is not enabled.", + description="Status codes from the quota status endpoint that mean quota tracking is unavailable rather than broken, such as a self-hosted deployment with rate limiting turned off. Every pool of the token whose request returned that status is then treated as untracked, so the authenticator stops waiting for quota resets, stops throttling proactively and stops rotating on exhaustion for it, while still signing requests. A token untracked this way stays untracked for the rest of the sync, because the endpoint is never consulted for it again, so a status the endpoint can also return transiently costs quota tracking for the whole run. If only some tokens return that status the others stay tracked, but they are no longer refreshed either, because the authenticator stops waiting for quota resets as soon as one token is untracked; once their counters are locally spent all traffic moves onto the untracked tokens. Rate limiting reported by ordinary responses is still handled by the stream's error handler, so one that retries 429 or 403 keeps working, and a retry rotates onto the next token; it pays the backoff the response asks for rather than the shortened one a tracked pool would get, since an untracked pool has no counters with which to argue the rejection was about that credential. Any status not listed still fails the connection, and this field never excuses a quota path missing from a response the endpoint did answer, so list only the codes the endpoint uses to report that rate limiting is not enabled. Do not list authentication or authorization statuses, since a 401 or 403 from a revoked credential would then be read as quota tracking being unavailable rather than as a credentials failure.", examples=[[404]], title="Unavailable Status Codes", + unique_items=True, ) 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 b36d4661e..0e63aec51 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4663,7 +4663,7 @@ def create_rate_limited_multiple_token_authenticator( "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 + # Normalize 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": [ @@ -4693,10 +4693,12 @@ def create_rate_limited_multiple_token_authenticator( key: str(InterpolatedString.create(value, parameters={}).eval(config)) for key, value in (model.quota_status_source.request_headers or {}).items() } - # Normalized the same way as the quota specs above, so an omitted field and an explicit - # `[]` key identically and keep sharing one set of counters. + # Normalize the same way as the quota specs above, so an omitted field and an explicit + # `[]` key identically and keep sharing one set of counters. Deduplicated as well as + # sorted, because the runtime turns this into a set: without it `[404]` and `[404, 404]` + # would key differently and stop sharing counters while behaving identically. quota_status_unavailable_status_codes = sorted( - model.quota_status_source.unavailable_status_codes or [] + set(model.quota_status_source.unavailable_status_codes or []) ) auth_method = model.auth_method or "Bearer" header = model.header or "Authorization" 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 461476746..bbce2a82a 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 @@ -11,6 +11,7 @@ import requests from pydantic.v1 import ValidationError +from airbyte_cdk.models import FailureType from airbyte_cdk.sources.declarative.auth.rate_limited_multiple_token import ( RateLimitedMultipleTokenAuthenticator, TokenQuota, @@ -259,13 +260,18 @@ def make_call(): assert authenticator._states[token]["rest"].remaining == 100 - used -def test_missing_path_in_quota_status_response_raises_config_error(requests_mock): +def test_missing_path_in_quota_status_response_raises_system_error(requests_mock): + """Reclassified from `config_error`: the quota paths come from the manifest, so a path the + response does not contain is a connector defect and there is nothing in the user's + configuration for them to correct.""" requests_mock.get(QUOTA_STATUS_URL, json={"unexpected": {}}) authenticator = _authenticator() - with pytest.raises(AirbyteTracedException, match="missing an expected field"): + with pytest.raises(AirbyteTracedException, match="does not contain the configured") as exc_info: authenticator(_prepared_request()) + assert exc_info.value.failure_type == FailureType.system_error + def test_no_tokens_raises_config_error(): with pytest.raises(AirbyteTracedException, match="tokens are missing"): @@ -1011,17 +1017,35 @@ def test_unavailable_status_is_untracked_and_never_blocks(requests_mock): ) authenticator = _authenticator(quota_status_unavailable_status_codes=[404]) - with patch("time.sleep") as sleep_mock: + with patch("time.sleep", side_effect=AssertionError("waited on an untracked quota")): for _ in range(3): - request = authenticator(_prepared_request()) + authenticator(_prepared_request()) - assert request.headers["Authorization"] == "token token_1" - sleep_mock.assert_not_called() for token in ("token_1", "token_2"): for pool in ("rest", "graphql"): assert authenticator._states[token][pool].tracked is False +def test_untracked_tokens_still_share_the_load(requests_mock): + """Every token hits the same `quota_status_url` and so gets the same status, which makes the + untracked branch the only one `_acquire_call` ever takes on a deployment that reports no + quota. Without advancing the active token there, one credential would serve the whole sync + and the rest of a multi-token configuration would go unused -- there is no counter saying a + token is spent, but there is also nothing saying the others should sit idle.""" + requests_mock.get( + QUOTA_STATUS_URL, status_code=404, json={"message": "Rate limiting is not enabled."} + ) + authenticator = _authenticator( + tokens=("token_1", "token_2", "token_3"), quota_status_unavailable_status_codes=[404] + ) + + used = [ + authenticator(_prepared_request()).headers["Authorization"].split()[1] for _ in range(9) + ] + + assert used == ["token_1", "token_2", "token_3"] * 3 + + def test_unavailable_status_without_opt_in_still_fails(requests_mock): """Unchanged behaviour for every connector that has not opted in -- the endpoint failing is still a broken connection, not a silent switch to untracked quotas.""" @@ -1084,41 +1108,211 @@ def test_untracked_pool_ignores_response_headers(requests_mock): ) request = _prepared_request() authenticator(request) + before = authenticator._states["token_1"]["rest"] + held_remaining, held_reset, held_limit = before.remaining, before.reset_at, before.limit response = requests.Response() response.status_code = 200 response.headers["X-RateLimit-Remaining"] = "17" response.headers["X-RateLimit-Reset"] = str(int(time.time()) + 3600) + response.headers["X-RateLimit-Limit"] = "5000" authenticator.update_from_response(request, response) - assert authenticator._states["token_1"]["rest"].tracked is False + state = authenticator._states["token_1"]["rest"] + assert state.tracked is False + # `tracked` is never written by `update_from_response`, so asserting only that would hold + # whether the guard exists or not. These three are what it protects. + assert (state.remaining, state.reset_at, state.limit) == ( + held_remaining, + held_reset, + held_limit, + ) -def test_missing_quota_path_is_untracked_only_for_that_pool(requests_mock): - """A deployment that reports some pools but not others should lose the pool, not the - connection -- and only when the connector has opted into tolerating this endpoint.""" +@pytest.mark.parametrize( + "unavailable_status_codes", + [pytest.param(None, id="without_opt_in"), pytest.param([404], id="with_opt_in")], +) +def test_missing_quota_path_always_raises(requests_mock, unavailable_status_codes): + """`unavailable_status_codes` says what an *error* from the endpoint means. It does not + excuse a path missing from a body the endpoint did answer with: a responding endpoint does + report quotas, so an absent path is a wrong path, and excusing it would let a typo in + `remaining_path` silently switch quota tracking off for the whole sync.""" body = _quota_status_body() del body["resources"]["graphql"] requests_mock.get(QUOTA_STATUS_URL, json=body) - authenticator = _authenticator(tokens=("token_1",), quota_status_unavailable_status_codes=[404]) + authenticator = _authenticator( + tokens=("token_1",), quota_status_unavailable_status_codes=unavailable_status_codes + ) - authenticator(_prepared_request()) + with pytest.raises(AirbyteTracedException) as exc_info: + authenticator(_prepared_request()) - assert authenticator._states["token_1"]["graphql"].tracked is False - assert authenticator._states["token_1"]["rest"].tracked is True - assert authenticator._states["token_1"]["rest"].remaining == 4999 + # The quota paths come from the manifest, not from anything the end user can edit, so there + # is no configuration for them to correct. + assert exc_info.value.failure_type == FailureType.system_error + assert "graphql" in exc_info.value.message -def test_missing_quota_path_without_opt_in_still_raises(requests_mock): - body = _quota_status_body() - del body["resources"]["graphql"] - requests_mock.get(QUOTA_STATUS_URL, json=body) - authenticator = _authenticator(tokens=("token_1",)) +def test_untracked_tokens_are_reported_once_with_the_right_scope(requests_mock): + """The consequence of untracking is only true of the tokens that are untracked. A message + claiming the connector will not wait, throttle or rotate -- while another token is still + tracked and doing all three -- points an operator at the wrong problem.""" + requests_mock.get( + QUOTA_STATUS_URL, + [ + {"status_code": 404, "json": {"message": "Rate limiting is not enabled."}}, + {"status_code": 200, "json": _quota_status_body()}, + ], + ) + authenticator = _authenticator(quota_status_unavailable_status_codes=[404]) - with pytest.raises(AirbyteTracedException): + with patch.object(authenticator._logger, "info") as info_mock: + authenticator(_prepared_request()) + + summaries = [ + call.args[0] % call.args[1:] if len(call.args) > 1 else call.args[0] + for call in info_mock.call_args_list + if "rate limiting is unavailable" in call.args[0] + ] + assert len(summaries) == 1, summaries + assert "for 1 of 2 tokens" in summaries[0] + # Not "the others are unaffected": once any token is untracked the exhaustion wait is + # unreachable, so the tracked token is never reseeded after its counters run out. + assert ( + "The other 1 keep proactive throttling until their counters are locally spent" + in (summaries[0]) + ) + assert "never refreshes them" in summaries[0] + + +def _mixed_authenticator(requests_mock, *, untracked_token, rest_remaining=5000): + """Seed one token from a 404 and the other from a healthy body. + + `_seed_all_tokens` fetches in `self._tokens` order, so the response list maps positionally + onto `token_1`, `token_2`. This is the only state in which four of the `tracked` guards are + reachable at all: every all-tracked or all-untracked run is short-circuited earlier by the + `_acquire_call` early return. + """ + unavailable = {"status_code": 404, "json": {"message": "Rate limiting is not enabled."}} + healthy = {"status_code": 200, "json": _quota_status_body(rest_remaining=rest_remaining)} + order = [unavailable, healthy] if untracked_token == "token_1" else [healthy, unavailable] + requests_mock.get(QUOTA_STATUS_URL, order) + authenticator = _authenticator(quota_status_unavailable_status_codes=[404]) + authenticator._ensure_initialized() + tracked_token = "token_2" if untracked_token == "token_1" else "token_1" + assert authenticator._states[untracked_token]["rest"].tracked is False + assert authenticator._states[tracked_token]["rest"].tracked is True + return authenticator + + +def test_exhausted_tracked_token_rotates_onto_an_untracked_token_without_waiting(requests_mock): + """An untracked token holds `remaining=0`, so a plain "every token is spent" test counts it + as exhausted and the connector sleeps for a reset that will never be reported. It should + rotate onto the untracked token instead, which can serve the request immediately.""" + authenticator = _mixed_authenticator(requests_mock, untracked_token="token_2") + for pool in ("rest", "graphql"): + authenticator._states["token_1"][pool].remaining = 0 + + with patch("time.sleep", side_effect=AssertionError("waited instead of rotating")): + request = authenticator(_prepared_request()) + + assert request.headers["Authorization"] == "token token_2" + + +def test_untracked_peer_disables_proactive_throttling(requests_mock): + """The budget delay is `seconds_until_reset / total_remaining` across every token. An + untracked token contributes 0 to the total and a reset that means nothing, so including it + invents a delay from a pool the server does not report.""" + authenticator = _mixed_authenticator( + requests_mock, untracked_token="token_2", rest_remaining=100 + ) + + with authenticator._lock: + assert authenticator._compute_budget_delay(authenticator._quotas[0]) is None + + with patch("time.sleep", side_effect=AssertionError("throttled an untracked pool")): authenticator(_prepared_request()) +def test_refresh_after_exhaustion_skips_the_reseed_when_a_token_is_untracked(requests_mock): + """Reachable under concurrency: a token can be untracked by another thread's reseed while + this one sleeps out the exhaustion wait. Reseeding again buys nothing, because + `_acquire_call` will rotate onto the untracked token rather than wait a second time.""" + authenticator = _mixed_authenticator(requests_mock, untracked_token="token_2") + authenticator._states["token_1"]["rest"].remaining = 0 + seeding_requests = requests_mock.call_count + + authenticator._refresh_after_exhaustion(authenticator._quotas[0]) + + assert requests_mock.call_count == seeding_requests + + +def test_untracked_sender_reports_no_alternative_token_even_when_another_token_has_quota( + requests_mock, +): + """`has_alternative_token` promises `HttpClient` that retrying in 0.1s will use a different + credential. `_acquire_call` returns the active token unchanged for an untracked pool, so an + untracked sender must answer False -- otherwise the retry hammers the credential the server + just rejected.""" + authenticator = _mixed_authenticator(requests_mock, untracked_token="token_1") + request = _prepared_request() + authenticator(request) + + assert request.headers["Authorization"] == "token token_1" + assert authenticator._states["token_2"]["rest"].remaining > 0 + assert authenticator.has_alternative_token(request) is False + + +def test_untracked_tokens_are_never_reseeded(requests_mock): + """ "Untracked holds for the rest of the sync" is load-bearing for the design, so pin the + mechanism rather than trusting the prose: the exhaustion wait is the only thing that reseeds + after startup, and it cannot fire while any token is untracked, so the quota endpoint is + never consulted again and an untracked pool can never silently flip back to tracked.""" + requests_mock.get( + QUOTA_STATUS_URL, + [ + {"status_code": 200, "json": _quota_status_body(rest_remaining=1, graphql_remaining=1)}, + {"status_code": 404, "json": {"message": "Rate limiting is not enabled."}}, + ], + ) + authenticator = _authenticator(quota_status_unavailable_status_codes=[404]) + authenticator._ensure_initialized() + seeding_requests = requests_mock.call_count + + # Spend the tracked token, then keep going well past the point where a reseed would happen + # if one were reachable. + with patch("time.sleep", side_effect=AssertionError("waited for a reset")): + for _ in range(6): + authenticator(_prepared_request()) + + assert requests_mock.call_count == seeding_requests + assert authenticator._states["token_2"]["rest"].tracked is False + + +def test_duplicate_unavailable_status_codes_are_rejected(): + """`[404, 404]` and `[404]` behave identically at runtime, so they must not be two different + manifests. The schema rejects the duplicate rather than silently collapsing it.""" + with pytest.raises(ValidationError): + RateLimitedMultipleTokenAuthenticatorModel( + type="RateLimitedMultipleTokenAuthenticator", + tokens=["token_1"], + quota_status_source={ + "type": "QuotaStatusSource", + "url": QUOTA_STATUS_URL, + "unavailable_status_codes": [404, 404], + }, + quotas=[ + { + "type": "TokenQuota", + "name": "rest", + "remaining_path": ["resources", "core", "remaining"], + "reset_path": ["resources", "core", "reset"], + } + ], + ) + + def test_unavailable_status_codes_are_threaded_through_the_factory(): """The manifest field has to reach the constructor, and two definitions that differ only by it must not collide in the factory's instance cache.""" @@ -1161,3 +1355,27 @@ def build(component_definition): strict = build(without) assert strict._unavailable_status_codes == set() assert strict is not tolerant, "the cache key must include the new field" + + # Order is not meaning: the runtime holds a set, so two definitions listing the same codes + # in a different order must keep sharing one set of counters. + reordered = build( + { + **definition, + "quota_status_source": { + "type": "QuotaStatusSource", + "url": QUOTA_STATUS_URL, + "unavailable_status_codes": [500, 404], + }, + } + ) + forward = build( + { + **definition, + "quota_status_source": { + "type": "QuotaStatusSource", + "url": QUOTA_STATUS_URL, + "unavailable_status_codes": [404, 500], + }, + } + ) + assert reordered is forward