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..064e7d6b7 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 @@ -61,6 +69,20 @@ 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`. 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. + """ class RateLimitedMultipleTokenAuthenticator(DeclarativeAuthenticator): @@ -104,6 +126,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 +157,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 +174,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 +198,89 @@ 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, + ) + 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". + + `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_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 + 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: return @@ -210,10 +320,30 @@ 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, 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 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 +392,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 @@ -291,11 +423,19 @@ 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( - 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() @@ -311,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] @@ -320,14 +461,22 @@ 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. `_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 ) @@ -338,14 +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: 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 @@ -383,6 +548,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 @@ -432,13 +603,24 @@ 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) 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 diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 091c694d2..e3e36ac56 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -444,6 +444,15 @@ 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 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: 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..944375e22 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -575,6 +575,13 @@ 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 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 a96c406d8..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,6 +4693,13 @@ 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() } + # 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( + set(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 +4730,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 +4766,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..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"): @@ -997,3 +1003,379 @@ 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", side_effect=AssertionError("waited on an untracked quota")): + for _ in range(3): + authenticator(_prepared_request()) + + 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.""" + 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) + 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) + + 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, + ) + + +@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=unavailable_status_codes + ) + + with pytest.raises(AirbyteTracedException) as exc_info: + authenticator(_prepared_request()) + + # 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_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 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.""" + 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" + + # 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