feat(low-code): let a quota status endpoint report that rate limiting is off

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

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