From 121fd81e89775d901fb0505a718f1595def220b9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 14:18:25 +0000 Subject: [PATCH 01/13] feat(low-code): add RateLimitedMultipleTokenAuthenticator declarative component Co-Authored-By: Daryna Ishchenko --- .../sources/declarative/auth/__init__.py | 11 +- .../auth/rate_limited_multiple_token.py | 288 ++++++++++++++++++ .../declarative_component_schema.yaml | 176 +++++++++++ .../models/declarative_component_schema.py | 116 +++++++ .../parsers/model_to_component_factory.py | 85 ++++++ .../auth/test_rate_limited_multiple_token.py | 265 ++++++++++++++++ 6 files changed, 940 insertions(+), 1 deletion(-) create mode 100644 airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py create mode 100644 unit_tests/sources/declarative/auth/test_rate_limited_multiple_token.py diff --git a/airbyte_cdk/sources/declarative/auth/__init__.py b/airbyte_cdk/sources/declarative/auth/__init__.py index 810437810d..df39801737 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 0000000000..454a2c3e77 --- /dev/null +++ b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py @@ -0,0 +1,288 @@ +# +# 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 + + 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: + 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) + self._acquire_call(quota) + request.headers.update(self.get_auth_header()) + 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 + return default_quota or self._quotas[0] + + def _acquire_call(self, quota: TokenQuota) -> None: + while True: + budget_delay: Optional[float] = None + wait_for_reset: Optional[float] = None + with self._lock: + state = self._states[self._active_token][quota.name] + if state.remaining > 0: + budget_delay = self._compute_budget_delay(quota) + state.remaining -= 1 + 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 = max(min_time_to_wait, 0) + 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 + + 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 7d99a01881..6fdb8f0b65 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,180 @@ 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" + examples: + - "PT2H" + - "PT30M" + 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 +2452,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 3bee30ca68..8f33ad7338 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 1e5801bb06..3382e1d9cc 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, @@ -403,9 +407,15 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( QueryProperties as QueryPropertiesModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + QuotaStatusSource as QuotaStatusSourceModel, +) 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, ) @@ -455,6 +465,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( SubstreamPartitionRouter as SubstreamPartitionRouterModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + TokenQuota as TokenQuotaModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( TypesMap as TypesMapModel, ) @@ -715,6 +728,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 +841,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 +4514,75 @@ 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(model.max_wait_time or "PT2H"), + 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/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 0000000000..b81ad9163c --- /dev/null +++ b/unit_tests/sources/declarative/auth/test_rate_limited_multiple_token.py @@ -0,0 +1,265 @@ +# +# 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 + + +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_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 From e31d33fb355887440991f7e5afda940a1027923a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 14:22:03 +0000 Subject: [PATCH 02/13] chore: remove unused model imports Co-Authored-By: Daryna Ishchenko --- .../declarative/parsers/model_to_component_factory.py | 6 ------ 1 file changed, 6 deletions(-) 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 3382e1d9cc..094f696c19 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -407,9 +407,6 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( QueryProperties as QueryPropertiesModel, ) -from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( - QuotaStatusSource as QuotaStatusSourceModel, -) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( Rate as RateModel, ) @@ -465,9 +462,6 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( SubstreamPartitionRouter as SubstreamPartitionRouterModel, ) -from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( - TokenQuota as TokenQuotaModel, -) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( TypesMap as TypesMapModel, ) From 51f1956f898d81cbb102d0a501f046b565e9bb98 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:42:15 +0000 Subject: [PATCH 03/13] fix: sign request with the token that was decremented; floor exhaustion wait Co-Authored-By: Daryna Ishchenko --- .../auth/rate_limited_multiple_token.py | 24 ++++++++++---- .../auth/test_rate_limited_multiple_token.py | 33 +++++++++++++++++++ 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py index 454a2c3e77..f85d6ddf9b 100644 --- a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py +++ b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py @@ -65,6 +65,7 @@ class RateLimitedMultipleTokenAuthenticator(DeclarativeAuthenticator): 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, @@ -126,14 +127,15 @@ def auth_header(self) -> str: @property def token(self) -> str: - return f"{self._auth_method} {self._active_token}".strip() + 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) - self._acquire_call(quota) - request.headers.update(self.get_auth_header()) + token = self._acquire_call(quota) + request.headers[self._header] = f"{self._auth_method} {token}".strip() return request def _ensure_initialized(self) -> None: @@ -152,14 +154,22 @@ def _match_quota(self, request: requests.PreparedRequest) -> TokenQuota: 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) -> None: + 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: - state = self._states[self._active_token][quota.name] + token = self._active_token + state = self._states[token][quota.name] if state.remaining > 0: budget_delay = self._compute_budget_delay(quota) state.remaining -= 1 @@ -176,7 +186,7 @@ def _acquire_call(self, quota: TokenQuota) -> None: 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 = max(min_time_to_wait, 0) + wait_for_reset = max(min_time_to_wait, self.MIN_EXHAUSTION_WAIT) else: self._active_token = next(self._tokens_iter) continue @@ -200,7 +210,7 @@ def _acquire_call(self, quota: TokenQuota) -> None: ) self._budget_logged = True time.sleep(budget_delay) - return + return token def _compute_budget_delay(self, quota: TokenQuota) -> Optional[float]: """Compute the proactive throttling delay. Must be called while holding the lock.""" 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 b81ad9163c..940e20db80 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 @@ -136,6 +136,11 @@ def test_waits_and_reseeds_when_all_tokens_exhausted(requests_mock): 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): @@ -178,6 +183,34 @@ def make_call(): 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": {}}) From 2e6fabf57d0e31ed7efef7fa950e54477f68c1b3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:49:11 +0000 Subject: [PATCH 04/13] fix: clamp exhaustion wait floor to max_wait_time Co-Authored-By: Daryna Ishchenko --- .../sources/declarative/auth/rate_limited_multiple_token.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 f85d6ddf9b..7cc854d7df 100644 --- a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py +++ b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py @@ -186,7 +186,10 @@ def _acquire_call(self, quota: TokenQuota) -> str: 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 = max(min_time_to_wait, self.MIN_EXHAUSTION_WAIT) + 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 From 77a2c801000848caa6aade7158317108f7ebd0c2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 12:20:57 +0000 Subject: [PATCH 05/13] fix: make Spec.generate_spec idempotent when advanced_auth enums were already converted Co-Authored-By: Daryna Ishchenko --- airbyte_cdk/sources/declarative/spec/spec.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/airbyte_cdk/sources/declarative/spec/spec.py b/airbyte_cdk/sources/declarative/spec/spec.py index 9d328023e7..5d91a5f0c8 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 From 5fd898f2c06c66e263998f7bcab015e1f0a036e1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:13:38 +0000 Subject: [PATCH 06/13] feat: interpolate max_wait_time on RateLimitedMultipleTokenAuthenticator Co-Authored-By: Daryna Ishchenko --- .../declarative/declarative_component_schema.yaml | 3 +++ .../parsers/model_to_component_factory.py | 8 +++++++- .../auth/test_rate_limited_multiple_token.py | 12 ++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 6fdb8f0b65..efdc7b4dcf 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -389,9 +389,12 @@ definitions: 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. 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 094f696c19..49af3be7ce 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4566,7 +4566,13 @@ def create_rate_limited_multiple_token_authenticator( quota_status_headers=quota_status_headers, auth_method=model.auth_method or "Bearer", header=model.header or "Authorization", - max_wait_time=parse_duration(model.max_wait_time or "PT2H"), + 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 ), 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 940e20db80..55eeb24c5e 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 @@ -296,3 +296,15 @@ def test_factory_returns_shared_instance_for_identical_definitions(): ) 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) From b66f3aa4991450b2c5dd7ca23eb839e82bc2f850 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:59:45 +0000 Subject: [PATCH 07/13] chore: retrigger CI Co-Authored-By: Daryna Ishchenko From f300579593c108453c454426097e011c94c2a699 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 15:06:38 +0000 Subject: [PATCH 08/13] fix: compute budget delay after decrementing quota counter Co-Authored-By: Daryna Ishchenko --- .../sources/declarative/auth/rate_limited_multiple_token.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 7cc854d7df..64236fdc85 100644 --- a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py +++ b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py @@ -171,8 +171,8 @@ def _acquire_call(self, quota: TokenQuota) -> str: token = self._active_token state = self._states[token][quota.name] if state.remaining > 0: - budget_delay = self._compute_budget_delay(quota) 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( ( From 891fb08ee27dc8503e8bb62d815aa7af7c1671b7 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:00:35 +0000 Subject: [PATCH 09/13] fix: exclude $parameters from shared-instance cache key Co-Authored-By: Daryna Ishchenko --- .../parsers/model_to_component_factory.py | 4 +++- .../auth/test_rate_limited_multiple_token.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) 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 807e1f3366..25abae3045 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4519,7 +4519,9 @@ def create_rate_limited_multiple_token_authenticator( ) -> 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) + # `$parameters` are excluded since they are propagated automatically and can + # differ per stream even when the authenticator definition is identical. + cache_key = model.json(exclude={"parameters"}, sort_keys=True) if cache_key in self._rate_limited_authenticators: return self._rate_limited_authenticators[cache_key] 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 55eeb24c5e..a5d6867c8b 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 @@ -298,6 +298,21 @@ def test_factory_returns_shared_instance_for_identical_definitions(): assert first is second +def test_factory_returns_shared_instance_when_only_parameters_differ(): + factory = ModelToComponentFactory() + config = {"pat": "token_1,token_2"} + + first_model = _model("{{ config['pat'] }}") + first_model.parameters = {"name": "stream_a"} + second_model = _model("{{ config['pat'] }}") + second_model.parameters = {"name": "stream_b"} + + first = factory.create_rate_limited_multiple_token_authenticator(first_model, config) + second = factory.create_rate_limited_multiple_token_authenticator(second_model, config) + + assert first is second + + def test_factory_interpolates_max_wait_time(): factory = ModelToComponentFactory() model = _model("{{ config['pat'] }}") From d92aaac4b7dc6552d643a8c9e7b3aabe237a271d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:11:06 +0000 Subject: [PATCH 10/13] fix: share rate-limited authenticator registry with substream factories; capture worker exceptions in thread-safety test Co-Authored-By: Daryna Ishchenko --- .../parsers/model_to_component_factory.py | 3 +++ .../auth/test_rate_limited_multiple_token.py | 12 +++++++++--- 2 files changed, 12 insertions(+), 3 deletions(-) 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 25abae3045..a65cae9d5f 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4142,6 +4142,9 @@ def create_parent_stream_config_with_substream_wrapper( ), api_budget=self._api_budget, ) + # Share the authenticator registry so parent and child streams draw from the + # same token quota counters + substream_factory._rate_limited_authenticators = self._rate_limited_authenticators return substream_factory.create_parent_stream_config( model=model, config=config, stream_name=stream_name, **kwargs 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 a5d6867c8b..eaecc22d77 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 @@ -193,12 +193,17 @@ def test_thread_safety_header_token_matches_decremented_token(requests_mock): ) calls = 150 # more than one token's quota, forcing rotation mid-flight signed_tokens = [] + worker_exceptions = [] signed_lock = threading.Lock() def make_call(): - request = authenticator(_prepared_request()) - with signed_lock: - signed_tokens.append(request.headers["Authorization"].split()[1]) + try: + request = authenticator(_prepared_request()) + with signed_lock: + signed_tokens.append(request.headers["Authorization"].split()[1]) + except Exception as exc: + with signed_lock: + worker_exceptions.append(exc) threads = [threading.Thread(target=make_call) for _ in range(calls)] for thread in threads: @@ -206,6 +211,7 @@ def make_call(): for thread in threads: thread.join() + assert worker_exceptions == [] for token in ("token_1", "token_2"): used = signed_tokens.count(token) assert authenticator._states[token]["rest"].remaining == 100 - used From 6c939e5b371223dc5db524026f2e59daccbd9aa5 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 20:48:48 +0000 Subject: [PATCH 11/13] fix: address deep review findings on rate-limited multiple token authenticator - Build shared-instance cache key from resolved constructor args so propagated $parameters cannot break instance sharing - Enforce cumulative max_wait_time across an exhaustion episode to prevent infinite reseed loops - Regenerate declarative_component_schema.py via codegen - Filter empty tokens in list form; validate max_wait_time is a fixed-length duration - Warn once on unmatched quota fallback; raise config_error for empty quotas - Share authenticator registry via constructor injection - Make Spec.generate_spec idempotent (normalize serialized dict, never mutate the model) Co-Authored-By: Daryna Ishchenko --- .../auth/rate_limited_multiple_token.py | 34 +- .../declarative_component_schema.yaml | 2 +- .../models/declarative_component_schema.py | 448 +++++++++--------- .../parsers/model_to_component_factory.py | 137 ++++-- airbyte_cdk/sources/declarative/spec/spec.py | 23 +- .../auth/test_rate_limited_multiple_token.py | 77 +++ .../sources/declarative/spec/test_spec.py | 18 + 7 files changed, 449 insertions(+), 290 deletions(-) diff --git a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py index 64236fdc85..377312f9da 100644 --- a/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py +++ b/airbyte_cdk/sources/declarative/auth/rate_limited_multiple_token.py @@ -87,8 +87,10 @@ def __init__( message="Authentication tokens are missing from the configuration.", ) if not quotas: - raise ValueError( - "RateLimitedMultipleTokenAuthenticator requires at least one quota pool" + raise AirbyteTracedException( + failure_type=FailureType.config_error, + internal_message="RateLimitedMultipleTokenAuthenticator requires at least one quota pool", + message="Quota pool configuration is missing.", ) self._logger = logging.getLogger("airbyte") self._tokens = list(tokens) @@ -106,6 +108,7 @@ def __init__( self._refresh_lock = threading.Lock() self._initialized = False self._budget_logged = False + self._unmatched_logged = False self._states: dict[str, dict[str, _QuotaState]] = {} self._token_to_http_client: Mapping[str, HttpClient] = { token: HttpClient( @@ -155,15 +158,22 @@ def _match_quota(self, request: requests.PreparedRequest) -> TokenQuota: 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, - ) + if not self._unmatched_logged: + self._logger.warning( + "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, + ) + self._unmatched_logged = True 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.""" + """Reserve one call from the matched pool and return the token it was charged to. + + `max_wait_time` bounds the *total* time spent waiting across all refresh attempts of a + single exhaustion episode, so stale reset timestamps cannot cause an endless reseed loop. + """ + exhaustion_deadline: Optional[float] = None while True: budget_delay: Optional[float] = None wait_for_reset: Optional[float] = None @@ -174,13 +184,17 @@ def _acquire_call(self, quota: TokenQuota) -> str: state.remaining -= 1 budget_delay = self._compute_budget_delay(quota) elif all(self._states[token][quota.name].remaining <= 0 for token in self._tokens): + now = time.monotonic() + if exhaustion_deadline is None: + exhaustion_deadline = now + self._max_wait_time.total_seconds() + remaining_budget = exhaustion_deadline - now 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(): + if remaining_budget <= 0 or min_time_to_wait >= remaining_budget: 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", @@ -188,7 +202,7 @@ def _acquire_call(self, quota: TokenQuota) -> str: ) wait_for_reset = min( max(min_time_to_wait, self.MIN_EXHAUSTION_WAIT), - self._max_wait_time.total_seconds(), + remaining_budget, ) else: self._active_token = next(self._tokens_iter) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index d62c53e55f..ec42d2d7a0 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -483,7 +483,7 @@ definitions: - ["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. + 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. Setting it on every pool is recommended so the reserve does not shrink when a sync starts with the pool already partially consumed. type: array items: type: string diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index f3c901945d..93e191bf40 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -6,7 +6,7 @@ from enum import Enum from typing import Any, Dict, List, Literal, Optional, Union -from pydantic.v1 import BaseModel, Extra, Field +from pydantic.v1 import BaseModel, Extra, Field, confloat, conint from airbyte_cdk.sources.declarative.models.base_model_with_deprecations import ( BaseModelWithDeprecations, @@ -18,12 +18,6 @@ class AuthFlowType(Enum): oauth1_0 = "oauth1.0" -class ScopesJoinStrategy(Enum): - space = "space" - comma = "comma" - plus = "plus" - - class BasicHttpAuthenticator(BaseModel): type: Literal["BasicHttpAuthenticator"] username: str = Field( @@ -52,15 +46,43 @@ class BearerAuthenticator(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class HttpMethod(Enum): + GET = "GET" + 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 DynamicStreamCheckConfig(BaseModel): type: Literal["DynamicStreamCheckConfig"] dynamic_stream_name: str = Field( ..., description="The dynamic stream name.", title="Dynamic Stream Name" ) - stream_count: Optional[int] = Field( + stream_count: Optional[conint(ge=1)] = Field( None, description="The number of streams to attempt reading from during a check operation. If unset, all generated streams are checked. Must be a positive integer; if it exceeds the total number of available streams, all streams are checked.", - ge=1, title="Stream Count", ) @@ -98,17 +120,16 @@ class ConcurrencyLevel(BaseModel): class ConstantBackoffStrategy(BaseModel): type: Literal["ConstantBackoffStrategy"] - backoff_time_in_seconds: Union[float, str] = Field( + backoff_time_in_seconds: Union[confloat(ge=0.0), str] = Field( ..., description="Backoff time in seconds.", examples=[30, 30.5, "{{ config['backoff_time'] }}"], title="Backoff Time", ) - jitter_range_in_seconds: Optional[float] = Field( + jitter_range_in_seconds: Optional[confloat(ge=0.0)] = Field( None, description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between backoff_time_in_seconds and backoff_time_in_seconds + (jitter_range_in_seconds * 2), so jitter only increases the base backoff.", examples=[15], - ge=0, title="Jitter Range", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -496,7 +517,7 @@ class Config: ) weight: Optional[Union[int, str]] = Field( None, - description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.", + description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.\n", title="Weight", ) @@ -516,6 +537,32 @@ class OnNoRecords(Enum): emit_parent = "emit_parent" +class RecordExpander(BaseModel): + type: Literal["RecordExpander"] + expand_records_from_field: List[str] = Field( + ..., + description="Path to a nested array field within each record. Items from this array will be extracted and emitted as separate records. Supports wildcards (*) for matching multiple arrays.", + examples=[ + ["lines", "data"], + ["items"], + ["nested", "array"], + ["sections", "*", "items"], + ], + title="Expand Records From Field", + ) + remain_original_record: Optional[bool] = Field( + False, + description='If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false.', + title="Remain Original Record", + ) + on_no_records: Optional[OnNoRecords] = Field( + OnNoRecords.skip, + description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', + title="On No Records", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class ExponentialBackoffStrategy(BaseModel): type: Literal["ExponentialBackoffStrategy"] factor: Optional[Union[float, str]] = Field( @@ -524,11 +571,10 @@ class ExponentialBackoffStrategy(BaseModel): examples=[5, 5.5, "10"], title="Factor", ) - jitter_range_in_seconds: Optional[float] = Field( + jitter_range_in_seconds: Optional[confloat(ge=0.0)] = Field( None, description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between computed_backoff and computed_backoff + (jitter_range_in_seconds * 2), so jitter only increases the computed backoff.", examples=[2], - ge=0, title="Jitter Range", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -549,125 +595,6 @@ class SessionTokenRequestBearerAuthenticator(BaseModel): type: Literal["Bearer"] -class HttpMethod(Enum): - GET = "GET" - 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" @@ -797,12 +724,13 @@ class JsonItemsDecoder(BaseModel): type: Literal["JsonItemsDecoder"] items_path: str = Field( ..., - description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax \u2014 do not include leading `$.` or trailing `[*]`.", + description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax — do not include leading `$.` or trailing `[*]`.", + examples=["dataByDepartmentAndSearchTerm", "dataByAsin", "data.users"], title="Items Path", ) encoding: Optional[str] = Field( "utf-8", - description="The character encoding of the JSON data. Defaults to UTF-8.", + description="Text encoding used to decode the streamed bytes before JSON parsing.", title="Encoding", ) @@ -965,22 +893,32 @@ class NoPagination(BaseModel): type: Literal["NoPagination"] -class State(BaseModel): +class Scope(BaseModel): class Config: extra = Extra.allow - min: int - max: int + scope: str = Field(..., description="The OAuth scope string to request from the provider.") -class OAuthScope(BaseModel): +class OptionalScope(BaseModel): class Config: extra = Extra.allow - scope: str = Field( - ..., - description="The OAuth scope string to request from the provider.", - ) + scope: str = Field(..., description="The OAuth scope string to request from the provider.") + + +class ScopesJoinStrategy(Enum): + space = "space" + comma = "comma" + plus = "plus" + + +class State(BaseModel): + class Config: + extra = Extra.allow + + min: int + max: int class OauthConnectorInputSpecification(BaseModel): @@ -1002,17 +940,13 @@ class Config: examples=["user:read user:read_orders workspaces:read"], title="Scopes", ) - # NOTE: scopes, optional_scopes, and scopes_join_strategy are processed by the - # platform OAuth handler (DeclarativeOAuthSpecHandler.kt), not by the CDK runtime. - # The CDK schema defines the manifest contract; the platform reads these fields - # during the OAuth consent flow to build the authorization URL. - scopes: Optional[List[OAuthScope]] = Field( + scopes: Optional[List[Scope]] = Field( None, description="List of OAuth scope objects. When present, takes precedence over the `scope` string property.\nThe scope values are joined using the `scopes_join_strategy` (default: space) before being\nsent to the OAuth provider.", examples=[[{"scope": "user:read"}, {"scope": "user:write"}]], title="Scopes", ) - optional_scopes: Optional[List[OAuthScope]] = Field( + optional_scopes: Optional[List[OptionalScope]] = Field( None, description="Optional OAuth scope objects that may or may not be granted.", examples=[[{"scope": "admin:read"}]], @@ -1385,7 +1319,14 @@ class AsyncJobStatusMap(BaseModel): completed: List[str] failed: List[str] timeout: List[str] - skipped: Optional[List[str]] = None + skipped: Optional[List[str]] = Field( + None, + description="Statuses that indicate the job was skipped because there is no data to return. Jobs with these statuses will not be retried and no records will be fetched.", + ) + + +class BlockSimultaneousSyncsAction(BaseModel): + type: Literal["BlockSimultaneousSyncsAction"] class ValueType(Enum): @@ -1747,6 +1688,40 @@ class AuthFlow(BaseModel): oauth_config_specification: Optional[OAuthConfigSpecification] = None +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. Setting it on every pool is recommended so the reserve does not shrink when a sync starts with the pool already partially consumed.", + 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 CheckStream(BaseModel): type: Literal["CheckStream"] stream_names: Optional[List[str]] = Field( @@ -2232,28 +2207,23 @@ class DefaultPaginator(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class RecordExpander(BaseModel): - type: Literal["RecordExpander"] - expand_records_from_field: List[str] = Field( +class DpathExtractor(BaseModel): + type: Literal["DpathExtractor"] + field_path: List[str] = Field( ..., - description="Path to a nested array field within each record. Items from this array will be extracted and emitted as separate records. Supports wildcards (*) for matching multiple arrays.", + description='List of potentially nested fields describing the full path of the field to extract. Use "*" to extract all values from an array. See more info in the [docs](https://docs.airbyte.com/connector-development/config-based/understanding-the-yaml-file/record-selector).', examples=[ - ["lines", "data"], - ["items"], - ["nested", "array"], - ["sections", "*", "items"], + ["data"], + ["data", "records"], + ["data", "{{ parameters.name }}"], + ["data", "*", "record"], ], - title="Expand Records From Field", - ) - remain_original_record: Optional[bool] = Field( - False, - description='If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false.', - title="Remain Original Record", + title="Field Path", ) - on_no_records: Optional[OnNoRecords] = Field( - OnNoRecords.skip, - description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', - title="On No Records", + record_expander: Optional[RecordExpander] = Field( + None, + description="Optional component to expand records by extracting items from nested array fields.", + title="Record Expander", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -2326,6 +2296,27 @@ class ListPartitionRouter(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class RecordSelector(BaseModel): + type: Literal["RecordSelector"] + extractor: Union[DpathExtractor, CustomRecordExtractor] + record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( + None, + description="Responsible for filtering records to be emitted by the Source.", + title="Record Filter", + ) + schema_normalization: Optional[Union[SchemaNormalization, CustomSchemaNormalization]] = Field( + None, + description="Responsible for normalization according to the schema.", + title="Schema Normalization", + ) + transform_before_filtering: Optional[bool] = Field( + None, + description="If true, transformation will be applied before record filtering.", + title="Transform Before Filtering", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class PaginationReset(BaseModel): type: Literal["PaginationReset"] action: Action1 @@ -2401,6 +2392,62 @@ class ConfigAddFields(BaseModel): ) +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", "PT{{ config.get('max_waiting_time', 120) }}M"], + 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 CompositeErrorHandler(BaseModel): type: Literal["CompositeErrorHandler"] error_handlers: List[Union[CompositeErrorHandler, DefaultErrorHandler, CustomErrorHandler]] = ( @@ -2446,27 +2493,6 @@ class Config: ) -class DpathExtractor(BaseModel): - type: Literal["DpathExtractor"] - field_path: List[str] = Field( - ..., - description='List of potentially nested fields describing the full path of the field to extract. Use "*" to extract all values from an array. See more info in the [docs](https://docs.airbyte.com/connector-development/config-based/understanding-the-yaml-file/record-selector).', - examples=[ - ["data"], - ["data", "records"], - ["data", "{{ parameters.name }}"], - ["data", "*", "record"], - ], - title="Field Path", - ) - record_expander: Optional[RecordExpander] = Field( - None, - description="Optional component to expand records by extracting items from nested array fields.", - title="Record Expander", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - class ZipfileDecoder(BaseModel): class Config: extra = Extra.allow @@ -2479,27 +2505,6 @@ class Config: ) -class RecordSelector(BaseModel): - type: Literal["RecordSelector"] - extractor: Union[DpathExtractor, CustomRecordExtractor] - record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( - None, - description="Responsible for filtering records to be emitted by the Source.", - title="Record Filter", - ) - schema_normalization: Optional[Union[SchemaNormalization, CustomSchemaNormalization]] = Field( - None, - description="Responsible for normalization according to the schema.", - title="Schema Normalization", - ) - transform_before_filtering: Optional[bool] = Field( - None, - description="If true, transformation will be applied before record filtering.", - title="Transform Before Filtering", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - class ConfigMigration(BaseModel): type: Literal["ConfigMigration"] description: Optional[str] = Field( @@ -2592,7 +2597,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -2632,7 +2637,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -3130,7 +3135,7 @@ class StateDelegatingStream(BaseModel): ) api_retention_period: Optional[str] = Field( None, - description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n * **PT1H**: 1 hour\n * **P1D**: 1 day\n * **P1W**: 1 week\n * **P1M**: 1 month\n * **P1Y**: 1 year\n * **P30D**: 30 days\n", + description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n* **PT1H**: 1 hour\n* **P1D**: 1 day\n* **P1W**: 1 week\n* **P1M**: 1 month\n* **P1Y**: 1 year\n* **P30D**: 30 days\n", examples=["P30D", "P90D", "P1Y"], title="API Retention Period", ) @@ -3230,10 +3235,9 @@ class AsyncRetriever(BaseModel): None, description="The time in minutes after which the single Async Job should be considered as Timed Out.", ) - failed_retry_wait_time_in_seconds: Optional[Union[int, str]] = Field( + failed_retry_wait_time_in_seconds: Optional[Union[conint(ge=1), str]] = Field( None, description="Time in seconds to wait before retrying a failed async job. Only applies to jobs that ran on the API side and reported a FAILED status (e.g. report generation failed due to a cooldown). Creation failures (HTTP errors when starting a job, such as 429s) and TIMED_OUT jobs are retried immediately and are not affected by this setting. When set, the orchestrator defers retry of real failed jobs until the wait time has elapsed, without blocking other jobs.", - ge=1, ) download_target_requester: Optional[Union[HttpRequester, CustomRequester]] = Field( None, @@ -3312,20 +3316,14 @@ class AsyncRetriever(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class BlockSimultaneousSyncsAction(BaseModel): - type: Literal["BlockSimultaneousSyncsAction"] - - class StreamGroup(BaseModel): - streams: List[str] = Field( + streams: List[DeclarativeStream] = Field( ..., - description='List of references to streams that belong to this group. Use JSON references to stream definitions (e.g., "#/definitions/my_stream").', + description="List of references to streams that belong to this group.\n", title="Streams", ) action: BlockSimultaneousSyncsAction = Field( - ..., - description="The action to apply to streams in this group.", - title="Action", + ..., description="The action to apply to streams in this group.", title="Action" ) 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 a65cae9d5f..a4b3e97fb2 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -7,6 +7,7 @@ import datetime import importlib import inspect +import json import logging import re from functools import partial @@ -707,6 +708,9 @@ def __init__( max_concurrent_async_job_count: Optional[int] = None, configured_catalog: Optional[ConfiguredAirbyteCatalog] = None, api_budget: Optional[APIBudget] = None, + rate_limited_authenticators: Optional[ + Dict[str, RateLimitedMultipleTokenAuthenticator] + ] = None, ): self._init_mappings() self._limit_pages_fetched_per_slice = limit_pages_fetched_per_slice @@ -723,7 +727,9 @@ 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._rate_limited_authenticators: Dict[str, RateLimitedMultipleTokenAuthenticator] = ( + rate_limited_authenticators if rate_limited_authenticators is not None else {} + ) self._job_tracker: JobTracker = JobTracker(max_concurrent_async_job_count or 1) # placeholder for deprecation warnings self._collected_deprecation_logs: List[ConnectorBuilderLogMessage] = [] @@ -4141,10 +4147,10 @@ def create_parent_stream_config_with_substream_wrapper( ), ), api_budget=self._api_budget, + # Share the authenticator registry so parent and child streams draw from the + # same token quota counters + rate_limited_authenticators=self._rate_limited_authenticators, ) - # Share the authenticator registry so parent and child streams draw from the - # same token quota counters - substream_factory._rate_limited_authenticators = self._rate_limited_authenticators return substream_factory.create_parent_stream_config( model=model, config=config, stream_name=stream_name, **kwargs @@ -4520,14 +4526,6 @@ def create_rate_limited_multiple_token_authenticator( 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). - # `$parameters` are excluded since they are propagated automatically and can - # differ per stream even when the authenticator definition is identical. - cache_key = model.json(exclude={"parameters"}, 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 "," @@ -4536,10 +4534,89 @@ def create_rate_limited_multiple_token_authenticator( ] else: tokens = [ - str(InterpolatedString.create(token, parameters={}).eval(config)) + token_value for token in model.tokens + if ( + token_value := str( + InterpolatedString.create(token, parameters={}).eval(config) + ).strip() + ) ] + quota_specs = [ + { + "name": quota_model.name, + "remaining_path": quota_model.remaining_path, + "reset_path": quota_model.reset_path, + "limit_path": quota_model.limit_path, + "matchers": [ + { + "method": matcher_model.method, + "url_base": matcher_model.url_base, + "url_path_pattern": matcher_model.url_path_pattern, + "params": matcher_model.params, + "headers": matcher_model.headers, + "weight": matcher_model.weight, + } + for matcher_model in quota_model.matchers or [] + ], + } + for quota_model in model.quotas + ] + + quota_status_url = str( + InterpolatedString.create(model.quota_status_source.url, parameters={}).eval(config) + ) + quota_status_http_method = ( + model.quota_status_source.http_method.value + if model.quota_status_source.http_method + else "GET" + ) + quota_status_headers = { + key: str(InterpolatedString.create(value, parameters={}).eval(config)) + for key, value in (model.quota_status_source.request_headers or {}).items() + } + auth_method = model.auth_method or "Bearer" + header = model.header or "Authorization" + max_wait_time_str = str( + InterpolatedString.create(model.max_wait_time or "PT2H", parameters={}).eval(config) + ) + max_wait_time = parse_duration(max_wait_time_str) + if not isinstance(max_wait_time, datetime.timedelta): + raise ValueError( + f"max_wait_time must be a fixed-length ISO 8601 duration (e.g. 'PT2H'); " + f"calendar-unit durations like '{max_wait_time_str}' are not supported" + ) + 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 + ) + + # Reuse the same instance for identical definitions so that all streams share the + # same token quota counters (similar to how api_budget is shared). The key is built + # from the resolved constructor arguments rather than the raw model so that + # stream-specific `$parameters` propagated onto the model (and its nested components) + # cannot break instance sharing. + cache_key = json.dumps( + { + "tokens": tokens, + "quotas": quota_specs, + "quota_status_url": quota_status_url, + "quota_status_http_method": quota_status_http_method, + "quota_status_headers": quota_status_headers, + "auth_method": auth_method, + "header": header, + "max_wait_time": max_wait_time.total_seconds(), + "budget_reserve_fraction": budget_reserve_fraction, + "budget_min_reserve": budget_min_reserve, + }, + sort_keys=True, + ) + if cache_key in self._rate_limited_authenticators: + return self._rate_limited_authenticators[cache_key] + quotas = [ TokenQuota( name=quota_model.name, @@ -4554,39 +4631,17 @@ def create_rate_limited_multiple_token_authenticator( 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_http_method=quota_status_http_method, 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 - ), + auth_method=auth_method, + header=header, + max_wait_time=max_wait_time, + budget_reserve_fraction=budget_reserve_fraction, + budget_min_reserve=budget_min_reserve, ) self._rate_limited_authenticators[cache_key] = authenticator return authenticator diff --git a/airbyte_cdk/sources/declarative/spec/spec.py b/airbyte_cdk/sources/declarative/spec/spec.py index 5d91a5f0c8..eaacfde812 100644 --- a/airbyte_cdk/sources/declarative/spec/spec.py +++ b/airbyte_cdk/sources/declarative/spec/spec.py @@ -55,20 +55,17 @@ def generate_spec(self) -> ConnectorSpecification: if self.documentation_url: obj["documentationUrl"] = self.documentation_url if self.advanced_auth: - 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: - oauth_input = getattr(oauth_spec, "oauth_connector_input_specification", None) - if ( - oauth_input - and hasattr(oauth_input, "scopes_join_strategy") - and isinstance(oauth_input.scopes_join_strategy, Enum) - ): - oauth_input.scopes_join_strategy = oauth_input.scopes_join_strategy.value # type: ignore + # Serialize to a dict and normalize enum values there so the typed model is + # never mutated and repeated calls produce the same result + advanced_auth = self.advanced_auth.dict() + if isinstance(advanced_auth.get("auth_flow_type"), Enum): + advanced_auth["auth_flow_type"] = advanced_auth["auth_flow_type"].value + oauth_spec = advanced_auth.get("oauth_config_specification") or {} + oauth_input = oauth_spec.get("oauth_connector_input_specification") or {} + if isinstance(oauth_input.get("scopes_join_strategy"), Enum): + oauth_input["scopes_join_strategy"] = oauth_input["scopes_join_strategy"].value # Map CDK AuthFlow model to protocol AdvancedAuth model - obj["advanced_auth"] = self.advanced_auth.dict() + obj["advanced_auth"] = advanced_auth # We remap these keys to camel case because that's the existing format expected by the rest of the platform return ConnectorSpecificationSerializer.load(obj) 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 eaecc22d77..077b08b79e 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 @@ -17,6 +17,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( RateLimitedMultipleTokenAuthenticator as RateLimitedMultipleTokenAuthenticatorModel, ) +from airbyte_cdk.sources.declarative.parsers.manifest_component_transformer import ( + ManifestComponentTransformer, +) from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import ( ModelToComponentFactory, ) @@ -232,6 +235,36 @@ def test_no_tokens_raises_config_error(): ) +def test_no_quotas_raises_config_error(): + with pytest.raises(AirbyteTracedException, match="Quota pool configuration is missing"): + RateLimitedMultipleTokenAuthenticator( + tokens=["token_1"], quotas=[], quota_status_url=QUOTA_STATUS_URL + ) + + +def test_raises_after_cumulative_max_wait_time(requests_mock): + """A quota status that keeps reporting zero remaining with a stale reset must not loop forever.""" + requests_mock.get( + QUOTA_STATUS_URL, + json=lambda request, context: _quota_status_body(rest_remaining=0, reset_in_seconds=-10), + ) + authenticator = _authenticator(max_wait_time=timedelta(seconds=30)) + clock = {"now": 0.0} + + def fake_sleep(seconds): + clock["now"] += seconds + + with ( + patch("time.sleep", side_effect=fake_sleep), + patch("time.monotonic", side_effect=lambda: clock["now"]), + ): + with pytest.raises(AirbyteTracedException, match="Rate limit is exceeded"): + authenticator(_prepared_request()) + + # total time waited never exceeds the configured max_wait_time + assert clock["now"] <= 30 + + def _model(tokens): return RateLimitedMultipleTokenAuthenticatorModel.parse_obj( { @@ -319,6 +352,50 @@ def test_factory_returns_shared_instance_when_only_parameters_differ(): assert first is second +def test_factory_returns_shared_instance_when_propagated_parameters_differ(): + """Streams with identical authenticator definitions but different propagated `$parameters` must share one instance.""" + factory = ModelToComponentFactory() + config = {"pat": "token_1,token_2"} + transformer = ManifestComponentTransformer() + + def _propagated_model(stream_name): + propagated = transformer.propagate_types_and_parameters( + "authenticator", + _model("{{ config['pat'] }}").dict(exclude_none=True, by_alias=True), + {"name": stream_name, "primary_key": "id"}, + ) + return RateLimitedMultipleTokenAuthenticatorModel.parse_obj(propagated) + + first_model = _propagated_model("stream_a") + second_model = _propagated_model("stream_b") + assert first_model.parameters != second_model.parameters + + first = factory.create_rate_limited_multiple_token_authenticator(first_model, config) + second = factory.create_rate_limited_multiple_token_authenticator(second_model, config) + + assert first is second + + +def test_factory_filters_empty_list_tokens(): + factory = ModelToComponentFactory() + model = _model(["{{ config['token_a'] }}", "{{ config['token_b'] }}", " "]) + + authenticator = factory.create_rate_limited_multiple_token_authenticator( + model, {"token_a": "token_1", "token_b": ""} + ) + + assert authenticator._tokens == ["token_1"] + + +def test_factory_rejects_calendar_unit_max_wait_time(): + factory = ModelToComponentFactory() + model = _model("{{ config['pat'] }}") + model.max_wait_time = "P1M" + + with pytest.raises(ValueError, match="calendar-unit"): + factory.create_rate_limited_multiple_token_authenticator(model, {"pat": "token_1"}) + + def test_factory_interpolates_max_wait_time(): factory = ModelToComponentFactory() model = _model("{{ config['pat'] }}") diff --git a/unit_tests/sources/declarative/spec/test_spec.py b/unit_tests/sources/declarative/spec/test_spec.py index 6516c146d1..9fab84d792 100644 --- a/unit_tests/sources/declarative/spec/test_spec.py +++ b/unit_tests/sources/declarative/spec/test_spec.py @@ -162,6 +162,24 @@ def test_spec(spec, expected_connection_specification) -> None: assert spec.generate_spec() == expected_connection_specification +def test_generate_spec_is_idempotent_and_does_not_mutate_the_model() -> None: + spec = component_spec( + connection_specification={"client_id": "my_client_id"}, + parameters={}, + advanced_auth=component_auth_flow( + auth_flow_type=component_auth_flow_type.oauth2_0, + predicate_key=None, + predicate_value=None, + ), + ) + + first = spec.generate_spec() + second = spec.generate_spec() + + assert first == second + assert spec.advanced_auth.auth_flow_type is component_auth_flow_type.oauth2_0 + + def test_given_list_of_transformations_when_transform_config_then_config_is_transformed() -> None: input_config = {"planet_code": "CRSC"} expected_config = { From 41be3c283cf1fc66d1ef0d3b3f827387507f4b39 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 20:52:21 +0000 Subject: [PATCH 12/13] fix: limit generated models diff to schema changes to keep mypy-compatible style Co-Authored-By: Daryna Ishchenko --- .../models/declarative_component_schema.py | 448 +++++++++--------- 1 file changed, 225 insertions(+), 223 deletions(-) diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 93e191bf40..b677f2b355 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -6,7 +6,7 @@ from enum import Enum from typing import Any, Dict, List, Literal, Optional, Union -from pydantic.v1 import BaseModel, Extra, Field, confloat, conint +from pydantic.v1 import BaseModel, Extra, Field from airbyte_cdk.sources.declarative.models.base_model_with_deprecations import ( BaseModelWithDeprecations, @@ -18,6 +18,12 @@ class AuthFlowType(Enum): oauth1_0 = "oauth1.0" +class ScopesJoinStrategy(Enum): + space = "space" + comma = "comma" + plus = "plus" + + class BasicHttpAuthenticator(BaseModel): type: Literal["BasicHttpAuthenticator"] username: str = Field( @@ -46,43 +52,15 @@ class BearerAuthenticator(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class HttpMethod(Enum): - GET = "GET" - 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 DynamicStreamCheckConfig(BaseModel): type: Literal["DynamicStreamCheckConfig"] dynamic_stream_name: str = Field( ..., description="The dynamic stream name.", title="Dynamic Stream Name" ) - stream_count: Optional[conint(ge=1)] = Field( + stream_count: Optional[int] = Field( None, description="The number of streams to attempt reading from during a check operation. If unset, all generated streams are checked. Must be a positive integer; if it exceeds the total number of available streams, all streams are checked.", + ge=1, title="Stream Count", ) @@ -120,16 +98,17 @@ class ConcurrencyLevel(BaseModel): class ConstantBackoffStrategy(BaseModel): type: Literal["ConstantBackoffStrategy"] - backoff_time_in_seconds: Union[confloat(ge=0.0), str] = Field( + backoff_time_in_seconds: Union[float, str] = Field( ..., description="Backoff time in seconds.", examples=[30, 30.5, "{{ config['backoff_time'] }}"], title="Backoff Time", ) - jitter_range_in_seconds: Optional[confloat(ge=0.0)] = Field( + jitter_range_in_seconds: Optional[float] = Field( None, description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between backoff_time_in_seconds and backoff_time_in_seconds + (jitter_range_in_seconds * 2), so jitter only increases the base backoff.", examples=[15], + ge=0, title="Jitter Range", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -517,7 +496,7 @@ class Config: ) weight: Optional[Union[int, str]] = Field( None, - description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.\n", + description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.", title="Weight", ) @@ -537,32 +516,6 @@ class OnNoRecords(Enum): emit_parent = "emit_parent" -class RecordExpander(BaseModel): - type: Literal["RecordExpander"] - expand_records_from_field: List[str] = Field( - ..., - description="Path to a nested array field within each record. Items from this array will be extracted and emitted as separate records. Supports wildcards (*) for matching multiple arrays.", - examples=[ - ["lines", "data"], - ["items"], - ["nested", "array"], - ["sections", "*", "items"], - ], - title="Expand Records From Field", - ) - remain_original_record: Optional[bool] = Field( - False, - description='If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false.', - title="Remain Original Record", - ) - on_no_records: Optional[OnNoRecords] = Field( - OnNoRecords.skip, - description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', - title="On No Records", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - class ExponentialBackoffStrategy(BaseModel): type: Literal["ExponentialBackoffStrategy"] factor: Optional[Union[float, str]] = Field( @@ -571,10 +524,11 @@ class ExponentialBackoffStrategy(BaseModel): examples=[5, 5.5, "10"], title="Factor", ) - jitter_range_in_seconds: Optional[confloat(ge=0.0)] = Field( + jitter_range_in_seconds: Optional[float] = Field( None, description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between computed_backoff and computed_backoff + (jitter_range_in_seconds * 2), so jitter only increases the computed backoff.", examples=[2], + ge=0, title="Jitter Range", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -595,6 +549,125 @@ class SessionTokenRequestBearerAuthenticator(BaseModel): type: Literal["Bearer"] +class HttpMethod(Enum): + GET = "GET" + 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. Setting it on every pool is recommended so the reserve does not shrink when a sync starts with the pool already partially consumed.", + 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" @@ -724,13 +797,12 @@ class JsonItemsDecoder(BaseModel): type: Literal["JsonItemsDecoder"] items_path: str = Field( ..., - description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax — do not include leading `$.` or trailing `[*]`.", - examples=["dataByDepartmentAndSearchTerm", "dataByAsin", "data.users"], + description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax \u2014 do not include leading `$.` or trailing `[*]`.", title="Items Path", ) encoding: Optional[str] = Field( "utf-8", - description="Text encoding used to decode the streamed bytes before JSON parsing.", + description="The character encoding of the JSON data. Defaults to UTF-8.", title="Encoding", ) @@ -893,32 +965,22 @@ class NoPagination(BaseModel): type: Literal["NoPagination"] -class Scope(BaseModel): - class Config: - extra = Extra.allow - - scope: str = Field(..., description="The OAuth scope string to request from the provider.") - - -class OptionalScope(BaseModel): +class State(BaseModel): class Config: extra = Extra.allow - scope: str = Field(..., description="The OAuth scope string to request from the provider.") - - -class ScopesJoinStrategy(Enum): - space = "space" - comma = "comma" - plus = "plus" + min: int + max: int -class State(BaseModel): +class OAuthScope(BaseModel): class Config: extra = Extra.allow - min: int - max: int + scope: str = Field( + ..., + description="The OAuth scope string to request from the provider.", + ) class OauthConnectorInputSpecification(BaseModel): @@ -940,13 +1002,17 @@ class Config: examples=["user:read user:read_orders workspaces:read"], title="Scopes", ) - scopes: Optional[List[Scope]] = Field( + # NOTE: scopes, optional_scopes, and scopes_join_strategy are processed by the + # platform OAuth handler (DeclarativeOAuthSpecHandler.kt), not by the CDK runtime. + # The CDK schema defines the manifest contract; the platform reads these fields + # during the OAuth consent flow to build the authorization URL. + scopes: Optional[List[OAuthScope]] = Field( None, description="List of OAuth scope objects. When present, takes precedence over the `scope` string property.\nThe scope values are joined using the `scopes_join_strategy` (default: space) before being\nsent to the OAuth provider.", examples=[[{"scope": "user:read"}, {"scope": "user:write"}]], title="Scopes", ) - optional_scopes: Optional[List[OptionalScope]] = Field( + optional_scopes: Optional[List[OAuthScope]] = Field( None, description="Optional OAuth scope objects that may or may not be granted.", examples=[[{"scope": "admin:read"}]], @@ -1319,14 +1385,7 @@ class AsyncJobStatusMap(BaseModel): completed: List[str] failed: List[str] timeout: List[str] - skipped: Optional[List[str]] = Field( - None, - description="Statuses that indicate the job was skipped because there is no data to return. Jobs with these statuses will not be retried and no records will be fetched.", - ) - - -class BlockSimultaneousSyncsAction(BaseModel): - type: Literal["BlockSimultaneousSyncsAction"] + skipped: Optional[List[str]] = None class ValueType(Enum): @@ -1688,40 +1747,6 @@ class AuthFlow(BaseModel): oauth_config_specification: Optional[OAuthConfigSpecification] = None -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. Setting it on every pool is recommended so the reserve does not shrink when a sync starts with the pool already partially consumed.", - 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 CheckStream(BaseModel): type: Literal["CheckStream"] stream_names: Optional[List[str]] = Field( @@ -2207,23 +2232,28 @@ class DefaultPaginator(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class DpathExtractor(BaseModel): - type: Literal["DpathExtractor"] - field_path: List[str] = Field( +class RecordExpander(BaseModel): + type: Literal["RecordExpander"] + expand_records_from_field: List[str] = Field( ..., - description='List of potentially nested fields describing the full path of the field to extract. Use "*" to extract all values from an array. See more info in the [docs](https://docs.airbyte.com/connector-development/config-based/understanding-the-yaml-file/record-selector).', + description="Path to a nested array field within each record. Items from this array will be extracted and emitted as separate records. Supports wildcards (*) for matching multiple arrays.", examples=[ - ["data"], - ["data", "records"], - ["data", "{{ parameters.name }}"], - ["data", "*", "record"], + ["lines", "data"], + ["items"], + ["nested", "array"], + ["sections", "*", "items"], ], - title="Field Path", + title="Expand Records From Field", ) - record_expander: Optional[RecordExpander] = Field( - None, - description="Optional component to expand records by extracting items from nested array fields.", - title="Record Expander", + remain_original_record: Optional[bool] = Field( + False, + description='If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false.', + title="Remain Original Record", + ) + on_no_records: Optional[OnNoRecords] = Field( + OnNoRecords.skip, + description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', + title="On No Records", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -2296,27 +2326,6 @@ class ListPartitionRouter(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class RecordSelector(BaseModel): - type: Literal["RecordSelector"] - extractor: Union[DpathExtractor, CustomRecordExtractor] - record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( - None, - description="Responsible for filtering records to be emitted by the Source.", - title="Record Filter", - ) - schema_normalization: Optional[Union[SchemaNormalization, CustomSchemaNormalization]] = Field( - None, - description="Responsible for normalization according to the schema.", - title="Schema Normalization", - ) - transform_before_filtering: Optional[bool] = Field( - None, - description="If true, transformation will be applied before record filtering.", - title="Transform Before Filtering", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - class PaginationReset(BaseModel): type: Literal["PaginationReset"] action: Action1 @@ -2392,62 +2401,6 @@ class ConfigAddFields(BaseModel): ) -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", "PT{{ config.get('max_waiting_time', 120) }}M"], - 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 CompositeErrorHandler(BaseModel): type: Literal["CompositeErrorHandler"] error_handlers: List[Union[CompositeErrorHandler, DefaultErrorHandler, CustomErrorHandler]] = ( @@ -2493,6 +2446,27 @@ class Config: ) +class DpathExtractor(BaseModel): + type: Literal["DpathExtractor"] + field_path: List[str] = Field( + ..., + description='List of potentially nested fields describing the full path of the field to extract. Use "*" to extract all values from an array. See more info in the [docs](https://docs.airbyte.com/connector-development/config-based/understanding-the-yaml-file/record-selector).', + examples=[ + ["data"], + ["data", "records"], + ["data", "{{ parameters.name }}"], + ["data", "*", "record"], + ], + title="Field Path", + ) + record_expander: Optional[RecordExpander] = Field( + None, + description="Optional component to expand records by extracting items from nested array fields.", + title="Record Expander", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class ZipfileDecoder(BaseModel): class Config: extra = Extra.allow @@ -2505,6 +2479,27 @@ class Config: ) +class RecordSelector(BaseModel): + type: Literal["RecordSelector"] + extractor: Union[DpathExtractor, CustomRecordExtractor] + record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( + None, + description="Responsible for filtering records to be emitted by the Source.", + title="Record Filter", + ) + schema_normalization: Optional[Union[SchemaNormalization, CustomSchemaNormalization]] = Field( + None, + description="Responsible for normalization according to the schema.", + title="Schema Normalization", + ) + transform_before_filtering: Optional[bool] = Field( + None, + description="If true, transformation will be applied before record filtering.", + title="Transform Before Filtering", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class ConfigMigration(BaseModel): type: Literal["ConfigMigration"] description: Optional[str] = Field( @@ -2597,7 +2592,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -2637,7 +2632,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -3135,7 +3130,7 @@ class StateDelegatingStream(BaseModel): ) api_retention_period: Optional[str] = Field( None, - description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n* **PT1H**: 1 hour\n* **P1D**: 1 day\n* **P1W**: 1 week\n* **P1M**: 1 month\n* **P1Y**: 1 year\n* **P30D**: 30 days\n", + description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n * **PT1H**: 1 hour\n * **P1D**: 1 day\n * **P1W**: 1 week\n * **P1M**: 1 month\n * **P1Y**: 1 year\n * **P30D**: 30 days\n", examples=["P30D", "P90D", "P1Y"], title="API Retention Period", ) @@ -3235,9 +3230,10 @@ class AsyncRetriever(BaseModel): None, description="The time in minutes after which the single Async Job should be considered as Timed Out.", ) - failed_retry_wait_time_in_seconds: Optional[Union[conint(ge=1), str]] = Field( + failed_retry_wait_time_in_seconds: Optional[Union[int, str]] = Field( None, description="Time in seconds to wait before retrying a failed async job. Only applies to jobs that ran on the API side and reported a FAILED status (e.g. report generation failed due to a cooldown). Creation failures (HTTP errors when starting a job, such as 429s) and TIMED_OUT jobs are retried immediately and are not affected by this setting. When set, the orchestrator defers retry of real failed jobs until the wait time has elapsed, without blocking other jobs.", + ge=1, ) download_target_requester: Optional[Union[HttpRequester, CustomRequester]] = Field( None, @@ -3316,14 +3312,20 @@ class AsyncRetriever(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class BlockSimultaneousSyncsAction(BaseModel): + type: Literal["BlockSimultaneousSyncsAction"] + + class StreamGroup(BaseModel): - streams: List[DeclarativeStream] = Field( + streams: List[str] = Field( ..., - description="List of references to streams that belong to this group.\n", + description='List of references to streams that belong to this group. Use JSON references to stream definitions (e.g., "#/definitions/my_stream").', title="Streams", ) action: BlockSimultaneousSyncsAction = Field( - ..., description="The action to apply to streams in this group.", title="Action" + ..., + description="The action to apply to streams in this group.", + title="Action", ) From dd2b20c5cc539f102c5ad3cef8ede649365bd36b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 14 Jul 2026 20:52:57 +0000 Subject: [PATCH 13/13] fix: add missing max_wait_time example to generated models Co-Authored-By: Daryna Ishchenko --- .../sources/declarative/models/declarative_component_schema.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index b677f2b355..9576c211af 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -652,7 +652,7 @@ class RateLimitedMultipleTokenAuthenticator(BaseModel): 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"], + examples=["PT2H", "PT30M", "PT{{ config.get('max_waiting_time', 120) }}M"], title="Maximum Wait Time", ) budget_reserve_fraction: Optional[float] = Field(