diff --git a/airbyte_cdk/sources/declarative/auth/__init__.py b/airbyte_cdk/sources/declarative/auth/__init__.py index 810437810..df3980173 100644 --- a/airbyte_cdk/sources/declarative/auth/__init__.py +++ b/airbyte_cdk/sources/declarative/auth/__init__.py @@ -4,5 +4,14 @@ from airbyte_cdk.sources.declarative.auth.jwt import JwtAuthenticator from airbyte_cdk.sources.declarative.auth.oauth import DeclarativeOauth2Authenticator +from airbyte_cdk.sources.declarative.auth.rate_limited_multiple_token import ( + RateLimitedMultipleTokenAuthenticator, + TokenQuota, +) -__all__ = ["DeclarativeOauth2Authenticator", "JwtAuthenticator"] +__all__ = [ + "DeclarativeOauth2Authenticator", + "JwtAuthenticator", + "RateLimitedMultipleTokenAuthenticator", + "TokenQuota", +] diff --git a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py new file mode 100644 index 000000000..64236fdc8 --- /dev/null +++ b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py @@ -0,0 +1,301 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import logging +import threading +import time +from dataclasses import dataclass, field +from datetime import timedelta +from itertools import cycle +from typing import Any, List, Mapping, Optional + +import requests + +from airbyte_cdk.models import FailureType +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.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 + + +@dataclass +class TokenQuota: + """A named per-token quota pool. + + `matchers` classify outgoing requests into the pool; a pool with no matchers acts as the + default pool. `remaining_path`/`reset_path`/`limit_path` locate the pool's values in the + quota status response. + """ + + name: str + remaining_path: List[str] + reset_path: List[str] + limit_path: Optional[List[str]] = None + matchers: List[RequestMatcher] = field(default_factory=list) + + +@dataclass +class _QuotaState: + remaining: int + reset_at: AirbyteDateTime + limit: int + + +class RateLimitedMultipleTokenAuthenticator(DeclarativeAuthenticator): + """Authenticator that rotates between multiple interchangeable tokens with per-token quota tracking. + + Each outgoing request is classified into a quota pool using the pool's request matchers. + The active token's counter for the matched pool is decremented locally; when it is exhausted + the authenticator rotates to the next token. When all tokens are exhausted for a pool, it + waits until the earliest quota reset (bounded by `max_wait_time`) and then refreshes all + counters from `quota_status_url`, or raises a transient error if the wait would be too long. + + A proactive throttling budget spreads the last calls over the time remaining until reset: + once every token's remaining count for a pool drops below its reserve + (`max(budget_min_reserve, budget_reserve_fraction * limit)`), a small delay proportional to + `seconds_until_reset / total_remaining` (capped at 10s) is injected before each request. + + Counters are seeded per token from `quota_status_url` on first use and refreshed after an + exhaustion wait. All state transitions are guarded by a lock so the authenticator can be + shared safely across concurrent streams; sleeps never hold the lock. + """ + + HEARTBEAT_INTERVAL = 60.0 # Log every 60s during exhaustion wait + MAX_BUDGET_DELAY = 10.0 # Cap for the per-request proactive throttling delay + MIN_EXHAUSTION_WAIT = 5.0 # Floor for the exhaustion wait, so stale reset timestamps can't cause a refresh busy-loop + + def __init__( + self, + tokens: List[str], + quotas: List[TokenQuota], + quota_status_url: str, + quota_status_http_method: str = "GET", + quota_status_headers: Optional[Mapping[str, str]] = None, + auth_method: str = "Bearer", + header: str = "Authorization", + max_wait_time: timedelta = timedelta(hours=2), + budget_reserve_fraction: float = 0.1, + budget_min_reserve: int = 50, + ) -> None: + if not tokens: + raise AirbyteTracedException( + failure_type=FailureType.config_error, + internal_message="RateLimitedMultipleTokenAuthenticator requires at least one token", + message="Authentication tokens are missing from the configuration.", + ) + if not quotas: + raise ValueError( + "RateLimitedMultipleTokenAuthenticator requires at least one quota pool" + ) + self._logger = logging.getLogger("airbyte") + self._tokens = list(tokens) + self._quotas = quotas + self._quota_status_url = quota_status_url + self._quota_status_http_method = quota_status_http_method + self._quota_status_headers = dict(quota_status_headers or {}) + self._auth_method = auth_method + self._header = header + self._max_wait_time = max_wait_time + self._budget_reserve_fraction = budget_reserve_fraction + self._budget_min_reserve = budget_min_reserve + + self._lock = threading.RLock() + self._refresh_lock = threading.Lock() + self._initialized = False + self._budget_logged = False + self._states: dict[str, dict[str, _QuotaState]] = {} + self._token_to_http_client: Mapping[str, HttpClient] = { + token: HttpClient( + name="quota_status", + logger=self._logger, + authenticator=TokenAuthenticator( + token, auth_method=self._auth_method, auth_header=self._header + ), + use_cache=False, # quota values change frequently; never reuse cached responses + ) + for token in self._tokens + } + self._tokens_iter = cycle(self._tokens) + self._active_token = next(self._tokens_iter) + + @property + def auth_header(self) -> str: + return self._header + + @property + def token(self) -> str: + with self._lock: + return f"{self._auth_method} {self._active_token}".strip() + + def __call__(self, request: requests.PreparedRequest) -> Any: + """Attach the HTTP headers required to authenticate on the HTTP request""" + self._ensure_initialized() + quota = self._match_quota(request) + token = self._acquire_call(quota) + request.headers[self._header] = f"{self._auth_method} {token}".strip() + return request + + def _ensure_initialized(self) -> None: + if self._initialized: + return + with self._refresh_lock: + if not self._initialized: + self._seed_all_tokens() + self._initialized = True + + def _match_quota(self, request: requests.PreparedRequest) -> TokenQuota: + default_quota: Optional[TokenQuota] = None + for quota in self._quotas: + if quota.matchers: + if any(matcher(request) for matcher in quota.matchers): + return quota + elif default_quota is None: + default_quota = quota + if default_quota is None: + self._logger.debug( + "Request %s did not match any quota pool; falling back to '%s'. Consider defining a matcher-less default pool.", + request.url, + self._quotas[0].name, + ) + return default_quota or self._quotas[0] + + def _acquire_call(self, quota: TokenQuota) -> str: + """Reserve one call from the matched pool and return the token it was charged to.""" + while True: + budget_delay: Optional[float] = None + wait_for_reset: Optional[float] = None + with self._lock: + token = self._active_token + state = self._states[token][quota.name] + 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): + min_time_to_wait = min( + ( + self._states[token][quota.name].reset_at - ab_datetime_now() + ).total_seconds() + for token in self._tokens + ) + if min_time_to_wait >= self._max_wait_time.total_seconds(): + raise AirbyteTracedException( + failure_type=FailureType.transient_error, + internal_message=f"Rate limits for all tokens (quota: {quota.name}) were reached and the next reset exceeds max_wait_time", + message="Rate limit is exceeded for all provided tokens.", + ) + wait_for_reset = min( + max(min_time_to_wait, self.MIN_EXHAUSTION_WAIT), + self._max_wait_time.total_seconds(), + ) + else: + self._active_token = next(self._tokens_iter) + continue + + if wait_for_reset is not None: + self._logger.info( + "All tokens exhausted (quota: %s). Waiting %.0fs until rate limit resets.", + quota.name, + wait_for_reset, + ) + self._sleep_with_heartbeat(wait_for_reset, quota.name) + self._refresh_after_exhaustion(quota) + continue + + if budget_delay is not None and budget_delay >= 0.1: + if not self._budget_logged: + self._logger.info( + "API budget: throttling requests (%.1fs delay) for quota '%s'.", + budget_delay, + quota.name, + ) + self._budget_logged = True + time.sleep(budget_delay) + return token + + 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 not all(state.remaining <= self._get_budget_reserve(state) for state in states): + return None + + active_state = self._states[self._active_token][quota.name] + seconds_to_reset = max((active_state.reset_at - ab_datetime_now()).total_seconds(), 0) + total_remaining = sum(max(state.remaining, 0) for state in states) + if total_remaining <= 0 or seconds_to_reset <= 0: + return None + + return min(seconds_to_reset / total_remaining, self.MAX_BUDGET_DELAY) + + def _get_budget_reserve(self, state: _QuotaState) -> float: + return max(self._budget_min_reserve, state.limit * self._budget_reserve_fraction) + + def _sleep_with_heartbeat(self, total_seconds: float, quota_name: str) -> None: + """Sleep for `total_seconds` in chunks, logging progress so operators can see the connector is not stuck.""" + remaining = total_seconds + while remaining > 0: + chunk = min(remaining, self.HEARTBEAT_INTERVAL) + time.sleep(chunk) + remaining -= chunk + if remaining > 0: + self._logger.info( + "Rate limit exhausted (quota: %s). Waiting for reset — %.0fs remaining.", + quota_name, + remaining, + ) + + def _refresh_after_exhaustion(self, quota: TokenQuota) -> None: + """Refresh counters after an exhaustion wait. Only one thread refreshes; others re-check state.""" + with self._refresh_lock: + with self._lock: + still_exhausted = all( + self._states[token][quota.name].remaining <= 0 for token in self._tokens + ) + if still_exhausted: + self._seed_all_tokens() + + def _seed_all_tokens(self) -> None: + states = {token: self._fetch_quota_states(token) for token in self._tokens} + with self._lock: + self._states = states + self._budget_logged = False + + def _fetch_quota_states(self, token: str) -> dict[str, _QuotaState]: + http_client = self._token_to_http_client[token] + _, response = http_client.send_request( + http_method=self._quota_status_http_method, + url=self._quota_status_url, + headers=self._quota_status_headers, + request_kwargs={}, + ) + 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) + limit = ( + self._extract_path(response_body, quota.limit_path) + if quota.limit_path + else remaining + ) + states[quota.name] = _QuotaState( + remaining=int(remaining), + reset_at=ab_datetime_parse(reset), + limit=int(limit), + ) + return states + + def _extract_path(self, response_body: Mapping[str, Any], path: List[str]) -> Any: + 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.", + ) + value = value[key] + return value diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 7d99a0188..efdc7b4dc 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -312,6 +312,7 @@ definitions: - "$ref": "#/definitions/LegacySessionTokenAuthenticator" - "$ref": "#/definitions/CustomAuthenticator" - "$ref": "#/definitions/NoAuth" + - "$ref": "#/definitions/RateLimitedMultipleTokenAuthenticator" examples: - authenticators: token: "#/definitions/ApiKeyAuthenticator" @@ -320,6 +321,183 @@ definitions: $parameters: type: object additionalProperties: true + RateLimitedMultipleTokenAuthenticator: + title: Rate Limited Multiple Token Authenticator + description: > + Authenticator that rotates between multiple interchangeable tokens, tracking each token's remaining + call quota. Outgoing requests are classified into quota pools using request matchers. When the active + token's quota for the matched pool is exhausted, the authenticator rotates to the next token. When all + tokens are exhausted, it waits until the earliest quota reset (bounded by `max_wait_time`) before + resuming. Quota counters are seeded per token from `quota_status_source` at startup and refreshed after + an exhaustion wait. + type: object + required: + - type + - tokens + - quota_status_source + - quotas + properties: + type: + type: string + enum: [RateLimitedMultipleTokenAuthenticator] + tokens: + title: Tokens + description: The tokens to rotate between. Either an explicit list of tokens, or a single string containing multiple tokens separated by `token_delimiter`. + anyOf: + - type: string + - type: array + items: + type: string + interpolation_context: + - config + examples: + - "{{ config['credentials']['personal_access_token'] }}" + - ["{{ config['token_1'] }}", "{{ config['token_2'] }}"] + token_delimiter: + title: Token Delimiter + description: Delimiter used to split a single token string into multiple tokens. + type: string + default: "," + auth_method: + title: Auth Method + description: "The prefix to prepend to the token in the auth header value (e.g. `Authorization: Bearer `)." + type: string + default: "Bearer" + examples: + - "Bearer" + - "token" + header: + title: Header Name + description: The name of the HTTP header in which to inject the token. + type: string + default: "Authorization" + quota_status_source: + title: Quota Status Source + description: Defines where to fetch each token's current quota status. Called once per token at startup and after an exhaustion wait, not per data request. + "$ref": "#/definitions/QuotaStatusSource" + quotas: + title: Quota Pools + description: > + Quota pools tracked per token. Each outgoing request is classified into the first pool whose + matchers match the request; a pool with no matchers acts as the default. The `remaining_path` and + `reset_path` locate each pool's values in the quota status response. + type: array + items: + "$ref": "#/definitions/TokenQuota" + max_wait_time: + title: Maximum Wait Time + description: ISO 8601 duration. When all tokens are exhausted, the maximum time to wait for a quota reset before raising a transient error. + type: string + default: "PT2H" + interpolation_context: + - config + examples: + - "PT2H" + - "PT30M" + - "PT{{ config.get('max_waiting_time', 120) }}M" + budget_reserve_fraction: + title: Budget Reserve Fraction + description: Fraction of each token's quota to keep in reserve. When every token drops below its reserve, requests are proactively throttled to spread the remaining calls until the quota reset. Set to 0 (along with `budget_min_reserve`) to disable throttling. + type: number + default: 0.1 + budget_min_reserve: + title: Budget Minimum Reserve + description: Minimum number of calls to keep in reserve per token before proactive throttling kicks in. + type: integer + default: 50 + $parameters: + type: object + additionalProperties: true + QuotaStatusSource: + title: Quota Status Source + description: Describes the endpoint from which a token's current rate limit quota status can be fetched. + type: object + required: + - type + - url + properties: + type: + type: string + enum: [QuotaStatusSource] + url: + title: URL + description: The full URL of the quota status endpoint. + type: string + interpolation_context: + - config + examples: + - "https://api.github.com/rate_limit" + - "{{ config.get('api_url', 'https://api.github.com') }}/rate_limit" + http_method: + title: HTTP Method + description: The HTTP method used to fetch the quota status. + type: string + enum: + - GET + - POST + default: GET + request_headers: + title: Request Headers + description: Additional headers to send with the quota status request. + type: object + additionalProperties: + type: string + $parameters: + type: object + additionalProperties: true + TokenQuota: + title: Token Quota + description: A named per-token quota pool with matchers that classify outgoing requests into the pool. + type: object + required: + - type + - name + - remaining_path + - reset_path + properties: + type: + type: string + enum: [TokenQuota] + name: + title: Name + description: Name of the quota pool. + type: string + examples: + - "rest" + - "graphql" + remaining_path: + title: Remaining Path + description: Path to the remaining call count for this pool in the quota status response. + type: array + items: + type: string + examples: + - ["resources", "core", "remaining"] + reset_path: + title: Reset Path + description: Path to the quota reset timestamp for this pool in the quota status response. + type: array + items: + type: string + examples: + - ["resources", "core", "reset"] + limit_path: + title: Limit Path + description: Optional path to the total call limit for this pool in the quota status response. Used to compute the proactive throttling reserve; falls back to the initially observed remaining count when not set. + type: array + items: + type: string + examples: + - ["resources", "core", "limit"] + matchers: + title: Matchers + description: List of matchers that classify outgoing requests into this quota pool. The first pool whose matcher matches a request is used. A pool with no matchers acts as the default pool. + type: array + items: + "$ref": "#/definitions/HttpRequestRegexMatcher" + $parameters: + type: object + additionalProperties: true CheckStream: title: Streams to Check description: Defines the streams to try reading when running a check operation. @@ -2277,6 +2455,7 @@ definitions: - "$ref": "#/definitions/CustomAuthenticator" - "$ref": "#/definitions/NoAuth" - "$ref": "#/definitions/LegacySessionTokenAuthenticator" + - "$ref": "#/definitions/RateLimitedMultipleTokenAuthenticator" fetch_properties_from_endpoint: deprecated: true deprecation_message: "Use `query_properties` field instead." diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 3bee30ca6..8f33ad733 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -549,6 +549,120 @@ class HttpMethod(Enum): POST = "POST" +class QuotaStatusSource(BaseModel): + type: Literal["QuotaStatusSource"] + url: str = Field( + ..., + description="The full URL of the quota status endpoint.", + examples=[ + "https://api.github.com/rate_limit", + "{{ config.get('api_url', 'https://api.github.com') }}/rate_limit", + ], + title="URL", + ) + http_method: Optional[HttpMethod] = Field( + HttpMethod.GET, + description="The HTTP method used to fetch the quota status.", + title="HTTP Method", + ) + request_headers: Optional[Dict[str, str]] = Field( + None, + description="Additional headers to send with the quota status request.", + title="Request Headers", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + +class TokenQuota(BaseModel): + type: Literal["TokenQuota"] + name: str = Field( + ..., + description="Name of the quota pool.", + examples=["rest", "graphql"], + title="Name", + ) + remaining_path: List[str] = Field( + ..., + description="Path to the remaining call count for this pool in the quota status response.", + examples=[["resources", "core", "remaining"]], + title="Remaining Path", + ) + reset_path: List[str] = Field( + ..., + description="Path to the quota reset timestamp for this pool in the quota status response.", + examples=[["resources", "core", "reset"]], + title="Reset Path", + ) + limit_path: Optional[List[str]] = Field( + None, + description="Optional path to the total call limit for this pool in the quota status response. Used to compute the proactive throttling reserve; falls back to the initially observed remaining count when not set.", + examples=[["resources", "core", "limit"]], + title="Limit Path", + ) + matchers: Optional[List[HttpRequestRegexMatcher]] = Field( + None, + description="List of matchers that classify outgoing requests into this quota pool. The first pool whose matcher matches a request is used. A pool with no matchers acts as the default pool.", + title="Matchers", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + +class RateLimitedMultipleTokenAuthenticator(BaseModel): + type: Literal["RateLimitedMultipleTokenAuthenticator"] + tokens: Union[str, List[str]] = Field( + ..., + description="The tokens to rotate between. Either an explicit list of tokens, or a single string containing multiple tokens separated by `token_delimiter`.", + examples=[ + "{{ config['credentials']['personal_access_token'] }}", + ["{{ config['token_1'] }}", "{{ config['token_2'] }}"], + ], + title="Tokens", + ) + token_delimiter: Optional[str] = Field( + ",", + description="Delimiter used to split a single token string into multiple tokens.", + title="Token Delimiter", + ) + auth_method: Optional[str] = Field( + "Bearer", + description="The prefix to prepend to the token in the auth header value (e.g. `Authorization: Bearer `).", + examples=["Bearer", "token"], + title="Auth Method", + ) + header: Optional[str] = Field( + "Authorization", + description="The name of the HTTP header in which to inject the token.", + title="Header Name", + ) + quota_status_source: QuotaStatusSource = Field( + ..., + description="Defines where to fetch each token's current quota status. Called once per token at startup and after an exhaustion wait, not per data request.", + title="Quota Status Source", + ) + quotas: List[TokenQuota] = Field( + ..., + description="Quota pools tracked per token. Each outgoing request is classified into the first pool whose matchers match the request; a pool with no matchers acts as the default. The `remaining_path` and `reset_path` locate each pool's values in the quota status response.\n", + title="Quota Pools", + ) + max_wait_time: Optional[str] = Field( + "PT2H", + description="ISO 8601 duration. When all tokens are exhausted, the maximum time to wait for a quota reset before raising a transient error.", + examples=["PT2H", "PT30M"], + title="Maximum Wait Time", + ) + budget_reserve_fraction: Optional[float] = Field( + 0.1, + description="Fraction of each token's quota to keep in reserve. When every token drops below its reserve, requests are proactively throttled to spread the remaining calls until the quota reset. Set to 0 (along with `budget_min_reserve`) to disable throttling.", + title="Budget Reserve Fraction", + ) + budget_min_reserve: Optional[int] = Field( + 50, + description="Minimum number of calls to keep in reserve per token before proactive throttling kicks in.", + title="Budget Minimum Reserve", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class Action(Enum): SUCCESS = "SUCCESS" FAIL = "FAIL" @@ -2566,6 +2680,7 @@ class Config: LegacySessionTokenAuthenticator, CustomAuthenticator, NoAuth, + RateLimitedMultipleTokenAuthenticator, ], ] = Field( ..., @@ -2795,6 +2910,7 @@ class HttpRequester(BaseModelWithDeprecations): CustomAuthenticator, NoAuth, LegacySessionTokenAuthenticator, + RateLimitedMultipleTokenAuthenticator, ] ] = Field( None, 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 1e5801bb0..49af3be7c 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -65,6 +65,10 @@ from airbyte_cdk.sources.declarative.auth.oauth import ( DeclarativeSingleUseRefreshTokenOauth2Authenticator, ) +from airbyte_cdk.sources.declarative.auth.rate_limited_multiple_token import ( + RateLimitedMultipleTokenAuthenticator, + TokenQuota, +) from airbyte_cdk.sources.declarative.auth.selective_authenticator import SelectiveAuthenticator from airbyte_cdk.sources.declarative.auth.token import ( ApiKeyAuthenticator, @@ -406,6 +410,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( Rate as RateModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + RateLimitedMultipleTokenAuthenticator as RateLimitedMultipleTokenAuthenticatorModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( RecordExpander as RecordExpanderModel, ) @@ -715,6 +722,8 @@ def __init__( ) self._connector_state_manager = connector_state_manager or ConnectorStateManager() self._api_budget: Optional[Union[APIBudget]] = api_budget + # Shared instances so all streams see the same token quota counters (like api_budget) + self._rate_limited_authenticators: Dict[str, RateLimitedMultipleTokenAuthenticator] = {} self._job_tracker: JobTracker = JobTracker(max_concurrent_async_job_count or 1) # placeholder for deprecation warnings self._collected_deprecation_logs: List[ConnectorBuilderLogMessage] = [] @@ -826,6 +835,7 @@ def _init_mappings(self) -> None: UnlimitedCallRatePolicyModel: self.create_unlimited_call_rate_policy, RateModel: self.create_rate, HttpRequestRegexMatcherModel: self.create_http_request_matcher, + RateLimitedMultipleTokenAuthenticatorModel: self.create_rate_limited_multiple_token_authenticator, GroupingPartitionRouterModel: self.create_grouping_partition_router, } @@ -4498,6 +4508,81 @@ def create_http_request_matcher( weight=weight, ) + def create_rate_limited_multiple_token_authenticator( + self, + model: RateLimitedMultipleTokenAuthenticatorModel, + config: Config, + **kwargs: Any, + ) -> RateLimitedMultipleTokenAuthenticator: + # Reuse the same instance for identical definitions so that all streams share + # the same token quota counters (similar to how api_budget is shared). + cache_key = model.json(sort_keys=True) + if cache_key in self._rate_limited_authenticators: + return self._rate_limited_authenticators[cache_key] + + if isinstance(model.tokens, str): + tokens_value = InterpolatedString.create(model.tokens, parameters={}).eval(config) + delimiter = model.token_delimiter or "," + tokens = [ + token.strip() for token in str(tokens_value).split(delimiter) if token.strip() + ] + else: + tokens = [ + str(InterpolatedString.create(token, parameters={}).eval(config)) + for token in model.tokens + ] + + quotas = [ + TokenQuota( + name=quota_model.name, + remaining_path=quota_model.remaining_path, + reset_path=quota_model.reset_path, + limit_path=quota_model.limit_path, + matchers=[ + self.create_http_request_matcher(matcher_model, config) + for matcher_model in quota_model.matchers or [] + ], + ) + for quota_model in model.quotas + ] + + quota_status_url = InterpolatedString.create( + model.quota_status_source.url, parameters={} + ).eval(config) + quota_status_headers = { + key: str(InterpolatedString.create(value, parameters={}).eval(config)) + for key, value in (model.quota_status_source.request_headers or {}).items() + } + + authenticator = RateLimitedMultipleTokenAuthenticator( + tokens=tokens, + quotas=quotas, + quota_status_url=quota_status_url, + quota_status_http_method=( + model.quota_status_source.http_method.value + if model.quota_status_source.http_method + else "GET" + ), + quota_status_headers=quota_status_headers, + auth_method=model.auth_method or "Bearer", + header=model.header or "Authorization", + max_wait_time=parse_duration( + str( + InterpolatedString.create(model.max_wait_time or "PT2H", parameters={}).eval( + config + ) + ) + ), + budget_reserve_fraction=( + model.budget_reserve_fraction if model.budget_reserve_fraction is not None else 0.1 + ), + budget_min_reserve=( + model.budget_min_reserve if model.budget_min_reserve is not None else 50 + ), + ) + self._rate_limited_authenticators[cache_key] = authenticator + return authenticator + def set_api_budget(self, component_definition: ComponentDefinition, config: Config) -> None: self._api_budget = self.create_component( model_type=HTTPAPIBudgetModel, component_definition=component_definition, config=config diff --git a/airbyte_cdk/sources/declarative/spec/spec.py b/airbyte_cdk/sources/declarative/spec/spec.py index 9d328023e..5d91a5f0c 100644 --- a/airbyte_cdk/sources/declarative/spec/spec.py +++ b/airbyte_cdk/sources/declarative/spec/spec.py @@ -3,6 +3,7 @@ # from dataclasses import InitVar, dataclass, field +from enum import Enum from typing import Any, List, Mapping, MutableMapping, Optional from airbyte_cdk.models import ( @@ -54,7 +55,8 @@ def generate_spec(self) -> ConnectorSpecification: if self.documentation_url: obj["documentationUrl"] = self.documentation_url if self.advanced_auth: - self.advanced_auth.auth_flow_type = self.advanced_auth.auth_flow_type.value # type: ignore # We know this is always assigned to an AuthFlow which has the auth_flow_type field + if isinstance(self.advanced_auth.auth_flow_type, Enum): + self.advanced_auth.auth_flow_type = self.advanced_auth.auth_flow_type.value # type: ignore # We know this is always assigned to an AuthFlow which has the auth_flow_type field # Convert scopes_join_strategy enum to its string value (same pattern as auth_flow_type above) oauth_spec = getattr(self.advanced_auth, "oauth_config_specification", None) if oauth_spec: @@ -62,7 +64,7 @@ def generate_spec(self) -> ConnectorSpecification: if ( oauth_input and hasattr(oauth_input, "scopes_join_strategy") - and oauth_input.scopes_join_strategy is not None + and isinstance(oauth_input.scopes_join_strategy, Enum) ): oauth_input.scopes_join_strategy = oauth_input.scopes_join_strategy.value # type: ignore # Map CDK AuthFlow model to protocol AdvancedAuth model 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 new file mode 100644 index 000000000..55eeb24c5 --- /dev/null +++ b/unit_tests/sources/declarative/auth/test_rate_limited_multiple_token.py @@ -0,0 +1,310 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import threading +import time +from datetime import timedelta +from unittest.mock import patch + +import pytest +import requests + +from airbyte_cdk.sources.declarative.auth.rate_limited_multiple_token import ( + RateLimitedMultipleTokenAuthenticator, + TokenQuota, +) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + RateLimitedMultipleTokenAuthenticator as RateLimitedMultipleTokenAuthenticatorModel, +) +from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import ( + ModelToComponentFactory, +) +from airbyte_cdk.sources.streams.call_rate import HttpRequestRegexMatcher +from airbyte_cdk.utils import AirbyteTracedException + +QUOTA_STATUS_URL = "https://api.example.com/rate_limit" + + +def _quota_status_body(rest_remaining=5000, graphql_remaining=5000, reset_in_seconds=3600): + reset = int(time.time()) + reset_in_seconds + return { + "resources": { + "core": {"remaining": rest_remaining, "reset": reset, "limit": 5000}, + "graphql": {"remaining": graphql_remaining, "reset": reset, "limit": 5000}, + } + } + + +def _quotas(): + return [ + TokenQuota( + name="rest", + remaining_path=["resources", "core", "remaining"], + reset_path=["resources", "core", "reset"], + limit_path=["resources", "core", "limit"], + ), + TokenQuota( + name="graphql", + remaining_path=["resources", "graphql", "remaining"], + reset_path=["resources", "graphql", "reset"], + limit_path=["resources", "graphql", "limit"], + matchers=[HttpRequestRegexMatcher(url_path_pattern="/graphql")], + ), + ] + + +def _authenticator(tokens=("token_1", "token_2"), **kwargs): + return RateLimitedMultipleTokenAuthenticator( + tokens=list(tokens), + quotas=_quotas(), + quota_status_url=QUOTA_STATUS_URL, + auth_method="token", + **kwargs, + ) + + +def _prepared_request(url="https://api.example.com/repos"): + return requests.Request(method="GET", url=url).prepare() + + +def test_seeding_and_header_injection(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _authenticator() + + request = authenticator(_prepared_request()) + + assert request.headers["Authorization"] == "token token_1" + # one seeding call per token + assert requests_mock.call_count == 2 + + +@pytest.mark.parametrize( + "url,expected_quota", + [ + pytest.param("https://api.example.com/repos", "rest", id="rest_request"), + pytest.param("https://api.example.com/graphql", "graphql", id="graphql_request"), + ], +) +def test_quota_matching(requests_mock, url, expected_quota): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _authenticator(tokens=("token_1",)) + + authenticator(_prepared_request(url)) + + states = authenticator._states["token_1"] + assert states[expected_quota].remaining == 4999 + other_quota = "graphql" if expected_quota == "rest" else "rest" + assert states[other_quota].remaining == 5000 + + +def test_rotation_when_active_token_exhausted(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _authenticator() + authenticator._ensure_initialized() + authenticator._states["token_1"]["rest"].remaining = 0 + + request = authenticator(_prepared_request()) + + assert request.headers["Authorization"] == "token token_2" + assert authenticator._states["token_2"]["rest"].remaining == 4999 + + +def test_raises_transient_error_when_reset_exceeds_max_wait_time(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body(rest_remaining=0)) + authenticator = _authenticator(max_wait_time=timedelta(seconds=10)) + + with pytest.raises(AirbyteTracedException, match="Rate limit is exceeded"): + authenticator(_prepared_request()) + + +def test_waits_and_reseeds_when_all_tokens_exhausted(requests_mock): + responses = iter( + [ + _quota_status_body(rest_remaining=0, reset_in_seconds=5), + _quota_status_body(rest_remaining=0, reset_in_seconds=5), + _quota_status_body(rest_remaining=5000), + _quota_status_body(rest_remaining=5000), + ] + ) + requests_mock.get(QUOTA_STATUS_URL, json=lambda request, context: next(responses)) + authenticator = _authenticator() + + with patch("time.sleep") as mock_sleep: + request = authenticator(_prepared_request()) + + assert mock_sleep.called + assert request.headers["Authorization"] == "token token_1" + assert authenticator._states["token_1"]["rest"].remaining == 4999 + # the exhaustion wait is floored so stale reset timestamps can't busy-loop the quota endpoint + assert ( + mock_sleep.call_args_list[0][0][0] + >= RateLimitedMultipleTokenAuthenticator.MIN_EXHAUSTION_WAIT + ) + + +def test_budget_throttling_delay_injected_when_all_tokens_below_reserve(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body(rest_remaining=100)) + authenticator = _authenticator() + + with patch("time.sleep") as mock_sleep: + authenticator(_prepared_request()) + + # reserve = max(50, 0.1 * 5000) = 500 > 100 remaining on both tokens -> throttled + assert mock_sleep.called + delay = mock_sleep.call_args[0][0] + assert 0.1 <= delay <= RateLimitedMultipleTokenAuthenticator.MAX_BUDGET_DELAY + + +def test_no_budget_throttling_when_tokens_have_headroom(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _authenticator() + + with patch("time.sleep") as mock_sleep: + authenticator(_prepared_request()) + + assert not mock_sleep.called + + +def test_thread_safety_no_lost_decrements(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json=_quota_status_body()) + authenticator = _authenticator(tokens=("token_1",)) + calls = 200 + + def make_call(): + authenticator(_prepared_request()) + + threads = [threading.Thread(target=make_call) for _ in range(calls)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert authenticator._states["token_1"]["rest"].remaining == 5000 - calls + + +def test_thread_safety_header_token_matches_decremented_token(requests_mock): + """Under concurrent rotation, each request must be signed with the token whose counter was decremented.""" + seed = _quota_status_body() + seed["resources"]["core"]["remaining"] = 100 + requests_mock.get(QUOTA_STATUS_URL, json=seed) + authenticator = _authenticator( + tokens=("token_1", "token_2"), budget_reserve_fraction=0, budget_min_reserve=0 + ) + calls = 150 # more than one token's quota, forcing rotation mid-flight + signed_tokens = [] + signed_lock = threading.Lock() + + def make_call(): + request = authenticator(_prepared_request()) + with signed_lock: + signed_tokens.append(request.headers["Authorization"].split()[1]) + + threads = [threading.Thread(target=make_call) for _ in range(calls)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + for token in ("token_1", "token_2"): + used = signed_tokens.count(token) + assert authenticator._states[token]["rest"].remaining == 100 - used + + +def test_missing_path_in_quota_status_response_raises_config_error(requests_mock): + requests_mock.get(QUOTA_STATUS_URL, json={"unexpected": {}}) + + authenticator = _authenticator() + with pytest.raises(AirbyteTracedException, match="missing an expected field"): + authenticator(_prepared_request()) + + +def test_no_tokens_raises_config_error(): + with pytest.raises(AirbyteTracedException, match="tokens are missing"): + RateLimitedMultipleTokenAuthenticator( + tokens=[], quotas=_quotas(), quota_status_url=QUOTA_STATUS_URL + ) + + +def _model(tokens): + return RateLimitedMultipleTokenAuthenticatorModel.parse_obj( + { + "type": "RateLimitedMultipleTokenAuthenticator", + "tokens": tokens, + "auth_method": "token", + "quota_status_source": { + "type": "QuotaStatusSource", + "url": "{{ config.get('api_url', 'https://api.example.com') }}/rate_limit", + }, + "quotas": [ + { + "type": "TokenQuota", + "name": "rest", + "remaining_path": ["resources", "core", "remaining"], + "reset_path": ["resources", "core", "reset"], + }, + { + "type": "TokenQuota", + "name": "graphql", + "remaining_path": ["resources", "graphql", "remaining"], + "reset_path": ["resources", "graphql", "reset"], + "matchers": [ + {"type": "HttpRequestRegexMatcher", "url_path_pattern": "/graphql"} + ], + }, + ], + } + ) + + +@pytest.mark.parametrize( + "tokens,config,expected_tokens", + [ + pytest.param( + "{{ config['pat'] }}", + {"pat": "token_1,token_2, token_3"}, + ["token_1", "token_2", "token_3"], + id="delimiter_separated_string", + ), + pytest.param( + ["{{ config['token_a'] }}", "{{ config['token_b'] }}"], + {"token_a": "token_1", "token_b": "token_2"}, + ["token_1", "token_2"], + id="explicit_list", + ), + ], +) +def test_factory_token_parsing(tokens, config, expected_tokens): + factory = ModelToComponentFactory() + + authenticator = factory.create_rate_limited_multiple_token_authenticator(_model(tokens), config) + + assert authenticator._tokens == expected_tokens + assert authenticator._quota_status_url == "https://api.example.com/rate_limit" + assert [quota.name for quota in authenticator._quotas] == ["rest", "graphql"] + + +def test_factory_returns_shared_instance_for_identical_definitions(): + factory = ModelToComponentFactory() + config = {"pat": "token_1,token_2"} + + first = factory.create_rate_limited_multiple_token_authenticator( + _model("{{ config['pat'] }}"), config + ) + second = factory.create_rate_limited_multiple_token_authenticator( + _model("{{ config['pat'] }}"), config + ) + + assert first is second + + +def test_factory_interpolates_max_wait_time(): + factory = ModelToComponentFactory() + model = _model("{{ config['pat'] }}") + model.max_wait_time = "PT{{ config.get('max_waiting_time', 120) }}M" + + authenticator = factory.create_rate_limited_multiple_token_authenticator( + model, {"pat": "token_1", "max_waiting_time": 30} + ) + + assert authenticator._max_wait_time == timedelta(minutes=30)