airbyte_cdk.sources.declarative.auth

 27@dataclass
 28class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAuthenticator):
 29    """
 30    Generates OAuth2.0 access tokens from an OAuth2.0 refresh token and client credentials based on
 31    a declarative connector configuration file. Credentials can be defined explicitly or via interpolation
 32    at runtime. The generated access token is attached to each request via the Authorization header.
 33
 34    Attributes:
 35        token_refresh_endpoint (Union[InterpolatedString, str]): The endpoint to refresh the access token
 36        client_id (Union[InterpolatedString, str]): The client id
 37        client_secret (Union[InterpolatedString, str]): Client secret (can be empty for APIs that support this)
 38        refresh_token (Union[InterpolatedString, str]): The token used to refresh the access token
 39        access_token_name (Union[InterpolatedString, str]): THe field to extract access token from in the response
 40        expires_in_name (Union[InterpolatedString, str]): The field to extract expires_in from in the response
 41        config (Mapping[str, Any]): The user-provided configuration as specified by the source's spec
 42        scopes (Optional[List[str]]): The scopes to request
 43        token_expiry_date (Optional[Union[InterpolatedString, str]]): The access token expiration date
 44        token_expiry_date_format str: format of the datetime; provide it if expires_in is returned in datetime instead of seconds
 45        token_expiry_is_time_of_expiration bool: set True it if expires_in is returned as time of expiration instead of the number seconds until expiration
 46        refresh_request_body (Optional[Mapping[str, Any]]): The request body to send in the refresh request
 47        refresh_request_headers (Optional[Mapping[str, Any]]): The request headers to send in the refresh request
 48        send_refresh_request_as_query_params (bool): When True, the standard refresh args (`grant_type`, `refresh_token`, client credentials when not in an `Authorization` header, scopes, plus any `refresh_request_body` extras) are sent on the URL query string instead of in the request body, and the body is emitted empty. Use this for OAuth providers like Gong that document their refresh endpoint with refresh args on the URL query string. Defaults to False.
 49        grant_type: The grant_type to request for access_token. If set to refresh_token, the refresh_token parameter has to be provided
 50        message_repository (MessageRepository): the message repository used to emit logs on HTTP requests
 51        refresh_token_error_status_codes (Tuple[int, ...]): Status codes to identify refresh token errors in response
 52        refresh_token_error_key (str): Key to identify refresh token error in response
 53        refresh_token_error_values (Tuple[str, ...]): List of values to check for exception during token refresh process
 54    """
 55
 56    config: Mapping[str, Any]
 57    parameters: InitVar[Mapping[str, Any]]
 58    client_id: Optional[Union[InterpolatedString, str]] = None
 59    client_secret: Optional[Union[InterpolatedString, str]] = None
 60    token_refresh_endpoint: Optional[Union[InterpolatedString, str]] = None
 61    refresh_token: Optional[Union[InterpolatedString, str]] = None
 62    scopes: Optional[List[str]] = None
 63    token_expiry_date: Optional[Union[InterpolatedString, str]] = None
 64    _token_expiry_date: Optional[AirbyteDateTime] = field(init=False, repr=False, default=None)
 65    token_expiry_date_format: Optional[str] = None
 66    token_expiry_is_time_of_expiration: bool = False
 67    access_token_name: Union[InterpolatedString, str] = "access_token"
 68    access_token_value: Optional[Union[InterpolatedString, str]] = None
 69    client_id_name: Union[InterpolatedString, str] = "client_id"
 70    client_secret_name: Union[InterpolatedString, str] = "client_secret"
 71    expires_in_name: Union[InterpolatedString, str] = "expires_in"
 72    refresh_token_name: Union[InterpolatedString, str] = "refresh_token"
 73    refresh_request_body: Optional[Mapping[str, Any]] = None
 74    refresh_request_headers: Optional[Mapping[str, Any]] = None
 75    send_refresh_request_as_query_params: bool = False
 76    grant_type_name: Union[InterpolatedString, str] = "grant_type"
 77    grant_type: Union[InterpolatedString, str] = "refresh_token"
 78    message_repository: MessageRepository = NoopMessageRepository()
 79    profile_assertion: Optional[DeclarativeAuthenticator] = None
 80    use_profile_assertion: Optional[Union[InterpolatedBoolean, str, bool]] = False
 81    refresh_token_error_status_codes: Tuple[int, ...] = ()
 82    refresh_token_error_key: str = ""
 83    refresh_token_error_values: Tuple[str, ...] = ()
 84
 85    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 86        super().__init__(
 87            refresh_token_error_status_codes=self.refresh_token_error_status_codes,
 88            refresh_token_error_key=self.refresh_token_error_key,
 89            refresh_token_error_values=self.refresh_token_error_values,
 90        )
 91        if self.token_refresh_endpoint is not None:
 92            self._token_refresh_endpoint: Optional[InterpolatedString] = InterpolatedString.create(
 93                self.token_refresh_endpoint, parameters=parameters
 94            )
 95        else:
 96            self._token_refresh_endpoint = None
 97        self._client_id_name = InterpolatedString.create(self.client_id_name, parameters=parameters)
 98        self._client_id = (
 99            InterpolatedString.create(self.client_id, parameters=parameters)
100            if self.client_id
101            else self.client_id
102        )
103        self._client_secret_name = InterpolatedString.create(
104            self.client_secret_name, parameters=parameters
105        )
106        self._client_secret = (
107            InterpolatedString.create(self.client_secret, parameters=parameters)
108            if self.client_secret
109            else self.client_secret
110        )
111        self._refresh_token_name = InterpolatedString.create(
112            self.refresh_token_name, parameters=parameters
113        )
114        if self.refresh_token is not None:
115            self._refresh_token: Optional[InterpolatedString] = InterpolatedString.create(
116                self.refresh_token, parameters=parameters
117            )
118        else:
119            self._refresh_token = None
120        self.access_token_name = InterpolatedString.create(
121            self.access_token_name, parameters=parameters
122        )
123        self.expires_in_name = InterpolatedString.create(
124            self.expires_in_name, parameters=parameters
125        )
126        self.grant_type_name = InterpolatedString.create(
127            self.grant_type_name, parameters=parameters
128        )
129        self.grant_type = InterpolatedString.create(
130            "urn:ietf:params:oauth:grant-type:jwt-bearer"
131            if self.use_profile_assertion
132            else self.grant_type,
133            parameters=parameters,
134        )
135        self._refresh_request_body = InterpolatedMapping(
136            self.refresh_request_body or {}, parameters=parameters
137        )
138        self._refresh_request_headers = InterpolatedMapping(
139            self.refresh_request_headers or {}, parameters=parameters
140        )
141        self._send_refresh_request_as_query_params = self.send_refresh_request_as_query_params
142        try:
143            if (
144                isinstance(self.token_expiry_date, (int, str))
145                and str(self.token_expiry_date).isdigit()
146            ):
147                self._token_expiry_date = ab_datetime_parse(self.token_expiry_date)
148            else:
149                self._token_expiry_date = (
150                    ab_datetime_parse(
151                        InterpolatedString.create(
152                            self.token_expiry_date, parameters=parameters
153                        ).eval(self.config)
154                    )
155                    if self.token_expiry_date
156                    else ab_datetime_now() - timedelta(days=1)
157                )
158        except ValueError as e:
159            raise ValueError(f"Invalid token expiry date format: {e}")
160        self.use_profile_assertion = (
161            InterpolatedBoolean(self.use_profile_assertion, parameters=parameters)
162            if isinstance(self.use_profile_assertion, str)
163            else self.use_profile_assertion
164        )
165        self.assertion_name = "assertion"
166
167        if self.access_token_value is not None:
168            self._access_token_value = InterpolatedString.create(
169                self.access_token_value, parameters=parameters
170            ).eval(self.config)
171        else:
172            self._access_token_value = None
173
174        self._access_token: Optional[str] = (
175            self._access_token_value if self.access_token_value else None
176        )
177
178        if not self.use_profile_assertion and any(
179            client_creds is None for client_creds in [self.client_id, self.client_secret]
180        ):
181            raise ValueError(
182                "OAuthAuthenticator configuration error: Both 'client_id' and 'client_secret' are required for the "
183                "basic OAuth flow."
184            )
185        if self.profile_assertion is None and self.use_profile_assertion:
186            raise ValueError(
187                "OAuthAuthenticator configuration error: 'profile_assertion' is required when using the profile assertion flow."
188            )
189        if self.get_grant_type() == "refresh_token" and self._refresh_token is None:
190            raise ValueError(
191                "OAuthAuthenticator configuration error: A 'refresh_token' is required when the 'grant_type' is set to 'refresh_token'."
192            )
193
194    def get_token_refresh_endpoint(self) -> Optional[str]:
195        if self._token_refresh_endpoint is not None:
196            refresh_token_endpoint: str = self._token_refresh_endpoint.eval(self.config)
197            if not refresh_token_endpoint:
198                raise ValueError(
199                    "OAuthAuthenticator was unable to evaluate token_refresh_endpoint parameter"
200                )
201            return refresh_token_endpoint
202        return None
203
204    def get_client_id_name(self) -> str:
205        return self._client_id_name.eval(self.config)  # type: ignore # eval returns a string in this context
206
207    def get_client_id(self) -> str:
208        client_id = self._client_id.eval(self.config) if self._client_id else self._client_id
209        if not client_id:
210            raise ValueError("OAuthAuthenticator was unable to evaluate client_id parameter")
211        return client_id  # type: ignore # value will be returned as a string, or an error will be raised
212
213    def get_client_secret_name(self) -> str:
214        return self._client_secret_name.eval(self.config)  # type: ignore # eval returns a string in this context
215
216    def get_client_secret(self) -> str:
217        client_secret = (
218            self._client_secret.eval(self.config) if self._client_secret else self._client_secret
219        )
220        if not client_secret:
221            # We've seen some APIs allowing empty client_secret so we will only log here
222            logger.warning(
223                "OAuthAuthenticator was unable to evaluate client_secret parameter hence it'll be empty"
224            )
225        return client_secret  # type: ignore # value will be returned as a string, which might be empty
226
227    def get_refresh_token_name(self) -> str:
228        return self._refresh_token_name.eval(self.config)  # type: ignore # eval returns a string in this context
229
230    def get_refresh_token(self) -> Optional[str]:
231        return None if self._refresh_token is None else str(self._refresh_token.eval(self.config))
232
233    def get_scopes(self) -> List[str]:
234        return self.scopes or []
235
236    def get_access_token_name(self) -> str:
237        return self.access_token_name.eval(self.config)  # type: ignore # eval returns a string in this context
238
239    def get_expires_in_name(self) -> str:
240        return self.expires_in_name.eval(self.config)  # type: ignore # eval returns a string in this context
241
242    def get_grant_type_name(self) -> str:
243        return self.grant_type_name.eval(self.config)  # type: ignore # eval returns a string in this context
244
245    def get_grant_type(self) -> str:
246        return self.grant_type.eval(self.config)  # type: ignore # eval returns a string in this context
247
248    def get_refresh_request_body(self) -> Mapping[str, Any]:
249        return self._refresh_request_body.eval(self.config)
250
251    def get_refresh_request_headers(self) -> Mapping[str, Any]:
252        return self._refresh_request_headers.eval(self.config)
253
254    def should_send_refresh_request_as_query_params(self) -> bool:
255        return self._send_refresh_request_as_query_params
256
257    def get_token_expiry_date(self) -> AirbyteDateTime:
258        if not self._has_access_token_been_initialized():
259            return AirbyteDateTime.from_datetime(datetime.min)
260        return self._token_expiry_date  # type: ignore # _token_expiry_date is an AirbyteDateTime. It is never None despite what mypy thinks
261
262    def _has_access_token_been_initialized(self) -> bool:
263        return self._access_token is not None
264
265    def set_token_expiry_date(self, value: AirbyteDateTime) -> None:
266        self._token_expiry_date = value
267
268    def get_assertion_name(self) -> str:
269        return self.assertion_name
270
271    def get_assertion(self) -> str:
272        if self.profile_assertion is None:
273            raise ValueError("profile_assertion is not set")
274        return self.profile_assertion.token
275
276    def build_refresh_request_body(self) -> Mapping[str, Any]:
277        """
278        Returns the request body to set on the refresh request
279
280        Override to define additional parameters
281        """
282        if self.use_profile_assertion:
283            return {
284                self.get_grant_type_name(): self.get_grant_type(),
285                self.get_assertion_name(): self.get_assertion(),
286            }
287        return super().build_refresh_request_body()
288
289    @property
290    def access_token(self) -> str:
291        if self._access_token is None:
292            raise ValueError("access_token is not set")
293        return self._access_token
294
295    @access_token.setter
296    def access_token(self, value: str) -> None:
297        self._access_token = value
298
299    @property
300    def _message_repository(self) -> MessageRepository:
301        """
302        Overriding AbstractOauth2Authenticator._message_repository to allow for HTTP request logs
303        """
304        return self.message_repository

 55@dataclass
 56class JwtAuthenticator(DeclarativeAuthenticator):
 57    """
 58    Generates a JSON Web Token (JWT) based on a declarative connector configuration file. The generated token is attached to each request via the Authorization header.
 59
 60    Attributes:
 61        config (Mapping[str, Any]): The user-provided configuration as specified by the source's spec
 62        secret_key (Union[InterpolatedString, str]): The secret key used to sign the JWT
 63        algorithm (Union[str, JwtAlgorithm]): The algorithm used to sign the JWT
 64        token_duration (Optional[int]): The duration in seconds for which the token is valid
 65        base64_encode_secret_key (Optional[Union[InterpolatedBoolean, str, bool]]): Whether to base64 encode the secret key
 66        header_prefix (Optional[Union[InterpolatedString, str]]): The prefix to add to the Authorization header
 67        kid (Optional[Union[InterpolatedString, str]]): The key identifier to be included in the JWT header
 68        typ (Optional[Union[InterpolatedString, str]]): The type of the JWT.
 69        cty (Optional[Union[InterpolatedString, str]]): The content type of the JWT.
 70        iss (Optional[Union[InterpolatedString, str]]): The issuer of the JWT.
 71        sub (Optional[Union[InterpolatedString, str]]): The subject of the JWT.
 72        aud (Optional[Union[InterpolatedString, str]]): The audience of the JWT.
 73        additional_jwt_headers (Optional[Mapping[str, Any]]): Additional headers to include in the JWT.
 74        additional_jwt_payload (Optional[Mapping[str, Any]]): Additional payload to include in the JWT.
 75    """
 76
 77    config: Mapping[str, Any]
 78    parameters: InitVar[Mapping[str, Any]]
 79    secret_key: Union[InterpolatedString, str]
 80    algorithm: Union[str, JwtAlgorithm]
 81    token_duration: Optional[int]
 82    base64_encode_secret_key: Optional[Union[InterpolatedBoolean, str, bool]] = False
 83    header_prefix: Optional[Union[InterpolatedString, str]] = None
 84    kid: Optional[Union[InterpolatedString, str]] = None
 85    typ: Optional[Union[InterpolatedString, str]] = None
 86    cty: Optional[Union[InterpolatedString, str]] = None
 87    iss: Optional[Union[InterpolatedString, str]] = None
 88    sub: Optional[Union[InterpolatedString, str]] = None
 89    aud: Optional[Union[InterpolatedString, str]] = None
 90    additional_jwt_headers: Optional[Mapping[str, Any]] = None
 91    additional_jwt_payload: Optional[Mapping[str, Any]] = None
 92    passphrase: Optional[Union[InterpolatedString, str]] = None
 93    request_option: Optional[RequestOption] = None
 94
 95    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 96        self._secret_key = InterpolatedString.create(self.secret_key, parameters=parameters)
 97        self._algorithm = (
 98            JwtAlgorithm(self.algorithm) if isinstance(self.algorithm, str) else self.algorithm
 99        )
100        self._base64_encode_secret_key = (
101            InterpolatedBoolean(self.base64_encode_secret_key, parameters=parameters)
102            if isinstance(self.base64_encode_secret_key, str)
103            else self.base64_encode_secret_key
104        )
105        self._token_duration = self.token_duration
106        self._header_prefix = (
107            InterpolatedString.create(self.header_prefix, parameters=parameters)
108            if self.header_prefix
109            else None
110        )
111        self._kid = InterpolatedString.create(self.kid, parameters=parameters) if self.kid else None
112        self._typ = InterpolatedString.create(self.typ, parameters=parameters) if self.typ else None
113        self._cty = InterpolatedString.create(self.cty, parameters=parameters) if self.cty else None
114        self._iss = InterpolatedString.create(self.iss, parameters=parameters) if self.iss else None
115        self._sub = InterpolatedString.create(self.sub, parameters=parameters) if self.sub else None
116        self._aud = InterpolatedString.create(self.aud, parameters=parameters) if self.aud else None
117        self._additional_jwt_headers = InterpolatedMapping(
118            self.additional_jwt_headers or {}, parameters=parameters
119        )
120        self._additional_jwt_payload = InterpolatedMapping(
121            self.additional_jwt_payload or {}, parameters=parameters
122        )
123        self._passphrase = (
124            InterpolatedString.create(self.passphrase, parameters=parameters)
125            if self.passphrase
126            else None
127        )
128
129        # When we first implemented the JWT authenticator, we assumed that the signed token was always supposed
130        # to be loaded into the request headers under the `Authorization` key. This is not always the case, but
131        # this default option allows for backwards compatibility to be retained for existing connectors
132        self._request_option = self.request_option or RequestOption(
133            inject_into=RequestOptionType.header, field_name="Authorization", parameters=parameters
134        )
135
136    def _get_jwt_headers(self) -> dict[str, Any]:
137        """
138        Builds and returns the headers used when signing the JWT.
139        """
140        headers = self._additional_jwt_headers.eval(self.config, json_loads=json.loads)
141        if any(prop in headers for prop in ["kid", "alg", "typ", "cty"]):
142            raise ValueError(
143                "'kid', 'alg', 'typ', 'cty' are reserved headers and should not be set as part of 'additional_jwt_headers'"
144            )
145
146        if self._kid:
147            headers["kid"] = self._kid.eval(self.config, json_loads=json.loads)
148        if self._typ:
149            headers["typ"] = self._typ.eval(self.config, json_loads=json.loads)
150        if self._cty:
151            headers["cty"] = self._cty.eval(self.config, json_loads=json.loads)
152        headers["alg"] = self._algorithm
153        return headers
154
155    def _get_jwt_payload(self) -> dict[str, Any]:
156        """
157        Builds and returns the payload used when signing the JWT.
158        """
159        now = int(datetime.now().timestamp())
160        exp = now + self._token_duration if isinstance(self._token_duration, int) else now
161        nbf = now
162
163        payload = self._additional_jwt_payload.eval(self.config, json_loads=json.loads)
164        if any(prop in payload for prop in ["iss", "sub", "aud", "iat", "exp", "nbf"]):
165            raise ValueError(
166                "'iss', 'sub', 'aud', 'iat', 'exp', 'nbf' are reserved properties and should not be set as part of 'additional_jwt_payload'"
167            )
168
169        if self._iss:
170            payload["iss"] = self._iss.eval(self.config, json_loads=json.loads)
171        if self._sub:
172            payload["sub"] = self._sub.eval(self.config, json_loads=json.loads)
173        if self._aud:
174            payload["aud"] = self._aud.eval(self.config, json_loads=json.loads)
175
176        payload["iat"] = now
177        payload["exp"] = exp
178        payload["nbf"] = nbf
179        return payload
180
181    def _get_secret_key(self) -> JwtKeyTypes:
182        """
183        Returns the secret key used to sign the JWT.
184        """
185        secret_key: str = self._secret_key.eval(self.config, json_loads=json.loads)
186
187        if self._passphrase:
188            passphrase_value = self._passphrase.eval(self.config, json_loads=json.loads)
189            if passphrase_value:
190                private_key = serialization.load_pem_private_key(
191                    secret_key.encode(),
192                    password=passphrase_value.encode(),
193                )
194                return cast(JwtKeyTypes, private_key)
195
196        return (
197            base64.b64encode(secret_key.encode()).decode()
198            if self._base64_encode_secret_key
199            else secret_key
200        )
201
202    def _get_signed_token(self) -> Union[str, Any]:
203        """
204        Signed the JWT using the provided secret key and algorithm and the generated headers and payload. For additional information on PyJWT see: https://pyjwt.readthedocs.io/en/stable/
205        """
206        try:
207            return jwt.encode(
208                payload=self._get_jwt_payload(),
209                key=self._get_secret_key(),
210                algorithm=self._algorithm,
211                headers=self._get_jwt_headers(),
212            )
213        except Exception as e:
214            raise ValueError(f"Failed to sign token: {e}")
215
216    def _get_header_prefix(self) -> Union[str, None]:
217        """
218        Returns the header prefix to be used when attaching the token to the request.
219        """
220        return (
221            self._header_prefix.eval(self.config, json_loads=json.loads)
222            if self._header_prefix
223            else None
224        )
225
226    @property
227    def auth_header(self) -> str:
228        options = self._get_request_options(RequestOptionType.header)
229        return next(iter(options.keys()), "")
230
231    @property
232    def token(self) -> str:
233        return (
234            f"{self._get_header_prefix()} {self._get_signed_token()}"
235            if self._get_header_prefix()
236            else self._get_signed_token()
237        )
238
239    def get_request_params(self) -> Mapping[str, Any]:
240        return self._get_request_options(RequestOptionType.request_parameter)
241
242    def get_request_body_data(self) -> Union[Mapping[str, Any], str]:
243        return self._get_request_options(RequestOptionType.body_data)
244
245    def get_request_body_json(self) -> Mapping[str, Any]:
246        return self._get_request_options(RequestOptionType.body_json)
247
248    def _get_request_options(self, option_type: RequestOptionType) -> Mapping[str, Any]:
249        options: MutableMapping[str, Any] = {}
250        if self._request_option.inject_into == option_type:
251            self._request_option.inject_into_request(options, self.token, self.config)
252        return options

 89class RateLimitedMultipleTokenAuthenticator(DeclarativeAuthenticator):
 90    """Authenticator that rotates between multiple interchangeable tokens with per-token quota tracking.
 91
 92    Each outgoing request is classified into a quota pool using the pool's request matchers.
 93    The active token's counter for the matched pool is decremented locally; when it is exhausted
 94    the authenticator rotates to the next token. When all tokens are exhausted for a pool, it
 95    waits until the earliest quota reset (bounded by `max_wait_time`) and then refreshes all
 96    counters from `quota_status_url`, or raises a transient error if the wait would be too long.
 97
 98    A proactive throttling budget spreads the last calls over the time remaining until reset:
 99    once every token's remaining count for a pool drops below its reserve
100    (`max(budget_min_reserve, budget_reserve_fraction * limit)`), a small delay proportional to
101    `seconds_until_reset / total_remaining` (capped at 10s) is injected before each request.
102
103    Implements `ResponseAwareAuthenticator` and `TokenRotatingAuthenticator` (see
104    `airbyte_cdk.sources.streams.http.requests_native_auth.protocols`), which is how `HttpClient`
105    feeds it responses and asks it whether a rate-limit wait can be skipped.
106
107    Counters are seeded per token from `quota_status_url` on first use and refreshed after an
108    exhaustion wait. When a pool declares response headers, `update_from_response` additionally
109    reconciles that pool against the server on every response, which keeps the counters honest
110    between seedings and makes the authenticator rotate off a token the server has rejected even
111    though the local count still looks healthy. All state transitions are guarded by a lock so
112    the authenticator can be shared safely across concurrent streams; sleeps never hold the lock.
113    """
114
115    HEARTBEAT_INTERVAL = 60.0  # Log every 60s during exhaustion wait
116    MAX_BUDGET_DELAY = 10.0  # Cap for the per-request proactive throttling delay
117    MIN_EXHAUSTION_WAIT = 5.0  # Floor for the exhaustion wait, so stale reset timestamps can't cause a refresh busy-loop
118    # How far behind the reset we hold a response's reset may be while still counting as the
119    # current window. Covers ordinary disagreement between the quota endpoint and response
120    # headers; anything older is treated as belonging to a window that has already rolled over.
121    RESET_SKEW_TOLERANCE = timedelta(seconds=60)
122
123    def __init__(
124        self,
125        tokens: List[str],
126        quotas: List[TokenQuota],
127        quota_status_url: str,
128        quota_status_http_method: str = "GET",
129        quota_status_headers: Optional[Mapping[str, str]] = None,
130        quota_status_unavailable_status_codes: Optional[List[int]] = None,
131        auth_method: str = "Bearer",
132        header: str = "Authorization",
133        max_wait_time: timedelta = timedelta(hours=2),
134        budget_reserve_fraction: float = 0.1,
135        budget_min_reserve: int = 50,
136    ) -> None:
137        if not tokens:
138            raise AirbyteTracedException(
139                failure_type=FailureType.config_error,
140                internal_message="RateLimitedMultipleTokenAuthenticator requires at least one token",
141                message="Authentication tokens are missing from the configuration.",
142            )
143        if not quotas:
144            raise AirbyteTracedException(
145                failure_type=FailureType.config_error,
146                internal_message="RateLimitedMultipleTokenAuthenticator requires at least one quota pool",
147                message="Quota pool configuration is missing.",
148            )
149        self._logger = logging.getLogger("airbyte")
150        self._tokens = list(tokens)
151        self._quotas = quotas
152        self._quota_status_url = quota_status_url
153        self._quota_status_http_method = quota_status_http_method
154        self._quota_status_headers = dict(quota_status_headers or {})
155        self._auth_method = auth_method
156        self._header = header
157        self._max_wait_time = max_wait_time
158        self._budget_reserve_fraction = budget_reserve_fraction
159        self._budget_min_reserve = budget_min_reserve
160
161        self._unavailable_status_codes = set(quota_status_unavailable_status_codes or [])
162
163        self._lock = threading.RLock()
164        self._refresh_lock = threading.Lock()
165        self._initialized = False
166        self._budget_logged = False
167        self._unmatched_logged = False
168        self._untracked_logged = False
169        self._states: dict[str, dict[str, _QuotaState]] = {}
170        self._token_to_http_client: Mapping[str, HttpClient] = {
171            token: HttpClient(
172                name="quota_status",
173                logger=self._logger,
174                authenticator=TokenAuthenticator(
175                    token, auth_method=self._auth_method, auth_header=self._header
176                ),
177                use_cache=False,  # quota values change frequently; never reuse cached responses
178                error_handler=self._quota_status_error_handler(),
179            )
180            for token in self._tokens
181        }
182        self._tokens_iter = cycle(self._tokens)
183        self._active_token = next(self._tokens_iter)
184
185    @property
186    def auth_header(self) -> str:
187        return self._header
188
189    @property
190    def token(self) -> str:
191        with self._lock:
192            return f"{self._auth_method} {self._active_token}".strip()
193
194    def __call__(self, request: requests.PreparedRequest) -> Any:
195        """Attach the HTTP headers required to authenticate on the HTTP request"""
196        self._ensure_initialized()
197        quota = self._match_quota(request)
198        token = self._acquire_call(quota)
199        request.headers[self._header] = f"{self._auth_method} {token}".strip()
200        return request
201
202    def _quota_status_error_handler(self) -> Optional[HttpStatusErrorHandler]:
203        """Error handling for the quota status request itself.
204
205        `None` keeps `HttpClient`'s default, under which every non-2xx fails the connection --
206        which is correct when the endpoint is expected to work. When the connector has declared
207        that some statuses mean "quota tracking is not enabled here", those are mapped to
208        `IGNORE` instead, so `send_request` hands the response back rather than raising and
209        `_fetch_quota_states` can decide what it means. Statuses outside the list keep failing.
210        """
211        if not self._unavailable_status_codes:
212            return None
213        return HttpStatusErrorHandler(
214            self._logger,
215            error_mapping={
216                **DEFAULT_ERROR_MAPPING,
217                **{
218                    status_code: ErrorResolution(
219                        response_action=ResponseAction.IGNORE,
220                        failure_type=FailureType.transient_error,
221                    )
222                    for status_code in self._unavailable_status_codes
223                },
224            },
225        )
226
227    def _untracked_states(self) -> dict[str, _QuotaState]:
228        """A state per pool meaning "the server tracks nothing here".
229
230        `remaining=0` is load-bearing rather than arbitrary: it is what keeps every
231        `remaining > 0` test in this class correct for an untracked pool without also having to
232        consult `tracked`. Nothing ever raises it, since `_acquire_call` only decrements and
233        `update_from_response` returns early for an untracked pool.
234        """
235        now = ab_datetime_now()
236        return {
237            quota.name: _QuotaState(remaining=0, reset_at=now, limit=0, tracked=False)
238            for quota in self._quotas
239        }
240
241    def _log_untracked_tokens(self, states: Mapping[str, Mapping[str, _QuotaState]]) -> None:
242        """Report untracked tokens once, scoped to how many of them there are.
243
244        Deliberately called with every token's states rather than from `_fetch_quota_states`,
245        which sees one token at a time. The consequence of untracking -- no exhaustion waits, no
246        proactive throttling, no rotation -- is only true of the tokens that are untracked, and
247        a per-token call site cannot know whether the others are. Claiming it globally while one
248        token is still tracked and still doing all three would send an operator looking for a
249        problem in the wrong place.
250        """
251        if self._untracked_logged:
252            return
253        untracked = [
254            token
255            for token, pools in states.items()
256            if any(not state.tracked for state in pools.values())
257        ]
258        if not untracked:
259            return
260        self._untracked_logged = True
261        if len(untracked) == len(self._tokens):
262            self._logger.info(
263                "Quota status endpoint reports that rate limiting is unavailable. Token quotas "
264                "are untracked: the connector will not wait for quota resets, throttle "
265                "proactively, or rotate tokens on exhaustion. Responses that report a rate "
266                "limit are still handled by the stream's error handler."
267            )
268        else:
269            # Not "the others are unaffected": `_acquire_call` rotates onto an untracked token
270            # rather than waiting, so the exhaustion wait -- and with it the only reseed after
271            # startup -- becomes unreachable as soon as one token is untracked. The tracked
272            # tokens keep throttling until their counters are locally spent and are then left
273            # spent for the rest of the sync.
274            self._logger.info(
275                "Quota status endpoint reports that rate limiting is unavailable for %d of %d "
276                "tokens. Those tokens are used without quota tracking. The other %d keep "
277                "proactive throttling until their counters are locally spent, after which "
278                "traffic moves onto the untracked tokens: the connector no longer waits for a "
279                "quota reset, so it never refreshes them.",
280                len(untracked),
281                len(states),
282                len(states) - len(untracked),
283            )
284
285    def _ensure_initialized(self) -> None:
286        if self._initialized:
287            return
288        with self._refresh_lock:
289            if not self._initialized:
290                self._seed_all_tokens()
291                self._initialized = True
292
293    def _match_quota(self, request: requests.PreparedRequest) -> TokenQuota:
294        default_quota: Optional[TokenQuota] = None
295        for quota in self._quotas:
296            if quota.matchers:
297                if any(matcher(request) for matcher in quota.matchers):
298                    return quota
299            elif default_quota is None:
300                default_quota = quota
301        if default_quota is None:
302            if not self._unmatched_logged:
303                self._logger.warning(
304                    "Request %s did not match any quota pool; falling back to '%s'. Consider defining a matcher-less default pool.",
305                    request.url,
306                    self._quotas[0].name,
307                )
308                self._unmatched_logged = True
309        return default_quota or self._quotas[0]
310
311    def _acquire_call(self, quota: TokenQuota) -> str:
312        """Reserve one call from the matched pool and return the token it was charged to.
313
314        `max_wait_time` bounds the *total* time spent waiting across all refresh attempts of a
315        single exhaustion episode, so stale reset timestamps cannot cause an endless reseed loop.
316        """
317        exhaustion_deadline: Optional[float] = None
318        while True:
319            budget_delay: Optional[float] = None
320            wait_for_reset: Optional[float] = None
321            with self._lock:
322                token = self._active_token
323                state = self._states[token][quota.name]
324                if not state.tracked:
325                    # Nothing to spend and nothing to wait for, but the tokens are still there
326                    # to spread load over. Every token hits the same `quota_status_url` and so
327                    # gets the same status, which means this branch is the *only* one taken on a
328                    # deployment that reports no quota -- so without advancing here, one
329                    # credential would serve the entire sync and the rest would go unused.
330                    # Round-robin is the right rule precisely because there are no counters:
331                    # nothing distinguishes the tokens, and the server may still enforce limits
332                    # it declines to report.
333                    #
334                    # Note the other half of the mechanism: once any token is untracked the
335                    # exhaustion branch below can never fire, so `_refresh_after_exhaustion` --
336                    # the only reseed after startup -- is unreachable, and a tracked token's
337                    # quota is never picked up again even after its window resets.
338                    self._active_token = next(self._tokens_iter)
339                    return token
340                if state.remaining > 0:
341                    state.remaining -= 1
342                    budget_delay = self._compute_budget_delay(quota)
343                elif all(
344                    self._states[token][quota.name].remaining <= 0
345                    and self._states[token][quota.name].tracked
346                    for token in self._tokens
347                ):
348                    now = time.monotonic()
349                    if exhaustion_deadline is None:
350                        exhaustion_deadline = now + self._max_wait_time.total_seconds()
351                    remaining_budget = exhaustion_deadline - now
352                    min_time_to_wait = min(
353                        (
354                            self._states[token][quota.name].reset_at - ab_datetime_now()
355                        ).total_seconds()
356                        for token in self._tokens
357                    )
358                    if remaining_budget <= 0 or min_time_to_wait >= remaining_budget:
359                        raise AirbyteTracedException(
360                            failure_type=FailureType.transient_error,
361                            internal_message=f"Rate limits for all tokens (quota: {quota.name}) were reached and the next reset exceeds max_wait_time",
362                            message="Rate limit is exceeded for all provided tokens.",
363                        )
364                    wait_for_reset = min(
365                        max(min_time_to_wait, self.MIN_EXHAUSTION_WAIT),
366                        remaining_budget,
367                    )
368                else:
369                    self._active_token = next(self._tokens_iter)
370                    continue
371
372            if wait_for_reset is not None:
373                self._logger.info(
374                    "All tokens exhausted (quota: %s). Waiting %.0fs until rate limit resets.",
375                    quota.name,
376                    wait_for_reset,
377                )
378                self._sleep_with_heartbeat(wait_for_reset, quota.name)
379                self._refresh_after_exhaustion(quota)
380                continue
381
382            if budget_delay is not None and budget_delay >= 0.1:
383                if not self._budget_logged:
384                    self._logger.info(
385                        "API budget: throttling requests (%.1fs delay) for quota '%s'.",
386                        budget_delay,
387                        quota.name,
388                    )
389                    self._budget_logged = True
390                time.sleep(budget_delay)
391            return token
392
393    def _compute_budget_delay(self, quota: TokenQuota) -> Optional[float]:
394        """Compute the proactive throttling delay. Must be called while holding the lock."""
395        states = [self._states[token][quota.name] for token in self._tokens]
396        if any(not state.tracked for state in states):
397            return None
398        if not all(state.remaining <= self._get_budget_reserve(state) for state in states):
399            return None
400
401        active_state = self._states[self._active_token][quota.name]
402        seconds_to_reset = max((active_state.reset_at - ab_datetime_now()).total_seconds(), 0)
403        total_remaining = sum(max(state.remaining, 0) for state in states)
404        if total_remaining <= 0 or seconds_to_reset <= 0:
405            return None
406
407        return min(seconds_to_reset / total_remaining, self.MAX_BUDGET_DELAY)
408
409    def _get_budget_reserve(self, state: _QuotaState) -> float:
410        return max(self._budget_min_reserve, state.limit * self._budget_reserve_fraction)
411
412    def _sleep_with_heartbeat(self, total_seconds: float, quota_name: str) -> None:
413        """Sleep for `total_seconds` in chunks, logging progress so operators can see the connector is not stuck."""
414        remaining = total_seconds
415        while remaining > 0:
416            chunk = min(remaining, self.HEARTBEAT_INTERVAL)
417            time.sleep(chunk)
418            remaining -= chunk
419            if remaining > 0:
420                self._logger.info(
421                    "Rate limit exhausted (quota: %s). Waiting for reset — %.0fs remaining.",
422                    quota_name,
423                    remaining,
424                )
425
426    def _refresh_after_exhaustion(self, quota: TokenQuota) -> None:
427        """Refresh counters after an exhaustion wait. Only one thread refreshes; others re-check state.
428
429        The `tracked` term is not reachable from a single-threaded run -- reaching the wait at
430        all requires every token to be tracked -- but it is reachable under concurrency, because
431        another thread's reseed can untrack a token while this one sleeps. Reseeding then buys
432        nothing: `_acquire_call` will rotate onto the untracked token instead of waiting again.
433        """
434        with self._refresh_lock:
435            with self._lock:
436                still_exhausted = all(
437                    self._states[token][quota.name].remaining <= 0
438                    and self._states[token][quota.name].tracked
439                    for token in self._tokens
440                )
441            if still_exhausted:
442                self._seed_all_tokens()
443
444    def _seed_all_tokens(self) -> None:
445        # The wholesale swap intentionally discards local decrements made by concurrent threads
446        # between the fetch and the swap: the server response is the closest thing to truth, and
447        # merging local decrements on top of it would double-count the calls the server has
448        # already observed. The worst case is a slight overcount of `remaining` (requests that
449        # were in flight during the fetch), which at most causes a few 429s near the quota
450        # boundary that the stream-level error handler absorbs.
451        states = {token: self._fetch_quota_states(token) for token in self._tokens}
452        with self._lock:
453            self._states = states
454            self._budget_logged = False
455        self._log_untracked_tokens(states)
456
457    def _fetch_quota_states(self, token: str) -> dict[str, _QuotaState]:
458        http_client = self._token_to_http_client[token]
459        _, response = http_client.send_request(
460            http_method=self._quota_status_http_method,
461            url=self._quota_status_url,
462            headers=self._quota_status_headers,
463            request_kwargs={},
464        )
465        if response.status_code in self._unavailable_status_codes:
466            # Only reachable when the connector opted in: without `unavailable_status_codes`
467            # the default error mapping raises before this point. `_seed_all_tokens` reports it
468            # once every token has been fetched, which is the first point at which the scope of
469            # the consequence is known.
470            return self._untracked_states()
471        response_body = response.json()
472
473        states = {}
474        for quota in self._quotas:
475            remaining = self._extract_path(
476                response_body, quota.remaining_path, quota.name, "remaining"
477            )
478            reset = self._extract_path(response_body, quota.reset_path, quota.name, "reset")
479            limit = (
480                self._extract_path(response_body, quota.limit_path, quota.name, "limit")
481                if quota.limit_path
482                else remaining
483            )
484            states[quota.name] = _QuotaState(
485                remaining=int(remaining),
486                reset_at=ab_datetime_parse(reset),
487                limit=int(limit),
488            )
489        return states
490
491    def _extract_path(
492        self, response_body: Mapping[str, Any], path: List[str], quota_name: str, field_name: str
493    ) -> Any:
494        """Read a configured quota path out of the response, or fail.
495
496        A path the response does not contain is a `system_error` rather than a `config_error`:
497        the paths come from the manifest, not from anything the end user can edit, so there is
498        no configuration for them to correct. `unavailable_status_codes` does not soften this --
499        it says what an endpoint answering with an error *means*, and an endpoint that answers
500        with a body does report quotas, so a path missing from that body is a wrong path.
501        """
502        value: Any = response_body
503        for key in path:
504            if not isinstance(value, Mapping) or key not in value:
505                raise AirbyteTracedException(
506                    failure_type=FailureType.system_error,
507                    internal_message=(
508                        f"Quota status response did not contain the {field_name} path {path} "
509                        f"configured for quota '{quota_name}'"
510                    ),
511                    message=(
512                        f"Quota status response does not contain the configured {field_name} "
513                        f'path for token quota "{quota_name}".'
514                    ),
515                )
516            value = value[key]
517        return value
518
519    def update_from_response(
520        self, request: requests.PreparedRequest, response: requests.Response
521    ) -> None:
522        """Reconcile the matched pool's counters against what the server reported.
523
524        Called once per HTTP attempt by `HttpClient`. Inert unless the matched pool declares
525        response headers, so behaviour is unchanged for pools configured only with paths.
526
527        The update is attributed to the token that *sent* the request rather than to the
528        currently active one: under concurrency the authenticator may have rotated between the
529        request going out and the response coming back.
530        """
531        quota = self._match_quota(request)
532        if not quota.is_response_aware:
533            return
534        token = self._token_from_request(request)
535        if token is None:
536            return
537
538        remaining = self._header_int(response, quota.remaining_header)
539        if remaining is None and response.status_code in quota.exhaustion_status_codes:
540            # The server rejected the call for rate-limit reasons but told us nothing about the
541            # remaining count. Treat the pool as spent so the next request rotates.
542            remaining = 0
543        reset_at = self._header_datetime(response, quota.reset_header)
544        limit = self._header_int(response, quota.limit_header)
545        if remaining is None and limit is None and reset_at is None:
546            return
547
548        with self._lock:
549            state = self._states.get(token, {}).get(quota.name)
550            if state is None:
551                return  # not seeded yet; the initial seeding is the more authoritative source
552            if not state.tracked:
553                # The quota status endpoint said this pool is not tracked. Response headers
554                # could contradict that, but adopting them would resurrect exhaustion waits and
555                # throttling on a deployment that has rate limiting switched off. Rate-limit
556                # *responses* remain the error handler's job either way.
557                return
558            if limit is not None and limit > 0 and (reset_at is None or reset_at >= state.reset_at):
559                # A response from an older window carries that window's limit. Taking it would
560                # skew the throttling reserve and, on a later reset-only response, refill the
561                # pool to the wrong capacity -- so the limit is accepted on the same terms as
562                # the reset itself. Assigned before the rollover branch below, which reads
563                # `state.limit` when a fresh window arrives with no remaining header.
564                state.limit = limit
565            if reset_at is not None and reset_at > state.reset_at:
566                # The quota window rolled over, so the local count is meaningless -- take the
567                # server's numbers wholesale. With no remaining header to go on, a fresh window
568                # is worth a full limit.
569                state.reset_at = reset_at
570                state.remaining = remaining if remaining is not None else state.limit
571            elif remaining is not None and (
572                (remaining <= 0 and response.status_code in quota.exhaustion_status_codes)
573                or reset_at is None
574                or reset_at >= state.reset_at - self.RESET_SKEW_TOLERANCE
575            ):
576                # Same window: only ever tighten the estimate. Responses arrive out of order and
577                # a slow one carries a stale, higher count, so handing those calls back would let
578                # concurrent requests overspend. The window itself is never moved backwards.
579                #
580                # A count from a window that has already rolled over describes a window that no
581                # longer exists, and `min` would pin the fresh pool to it for the rest of the
582                # hour, so those are ignored -- with two exceptions.
583                #
584                # First, a zero on a response the pool counts as rate-limited is an exhaustion
585                # signal and must never be dropped, or a rate limit whose reset header trails the
586                # value we hold would silently stop rotation. The status check is what keeps that
587                # narrow: a zero on a *successful* response is just the last call of a window, and
588                # honouring it from a dead window would park a pool that has already refilled.
589                #
590                # Second, a reset within `RESET_SKEW_TOLERANCE` is treated as the current window,
591                # since the quota endpoint and the response headers can disagree by a little.
592                state.remaining = min(state.remaining, remaining)
593
594    def has_alternative_token(self, request: requests.PreparedRequest) -> bool:
595        """Whether another token could serve this request right now.
596
597        Answers the question a rate-limit backoff cannot answer for itself: the wait computed
598        from a response's reset header assumes the only way forward is for that quota to come
599        back, which is false when a different token still has calls. `HttpClient` uses this to
600        retry promptly instead of sleeping out a window it does not need.
601
602        Deliberately narrow. It reports True only when the token that *sent* the request is
603        spent for the matched pool -- so the next request is guaranteed to rotate -- and some
604        other token is not. If the sending token still has calls locally, the rejection was not
605        about exhausting it (a secondary limit, say, which on many APIs is per-user and would
606        reject every token alike), and waiting remains the right response.
607
608        An untracked sender answers False too, but for a different reason, and it is a trade-off
609        rather than a clear win. The retry does rotate -- `_acquire_call` round-robins untracked
610        tokens -- so what this withholds is only the *skipped wait*. The backoff it would skip is
611        computed from what the server said (a reset or `Retry-After` header), and an untracked
612        pool has no counters with which to argue the rejection was about this credential
613        specifically. Overriding the server's own instruction on a guess would, when the limit is
614        shared across credentials, burn every retry in under a second and fail a request that
615        waiting would have completed. So a rate-limited response on an untracked pool rotates
616        credentials but still pays the computed backoff.
617        """
618        quota = self._match_quota(request)
619        sender = self._token_from_request(request)
620        with self._lock:
621            if not self._states or sender is None:
622                return False
623            sender_state = self._states[sender][quota.name]
624            if not sender_state.tracked or sender_state.remaining > 0:
625                return False
626            return any(
627                self._states[token][quota.name].remaining > 0
628                for token in self._tokens
629                if token != sender
630            )
631
632    def _token_from_request(self, request: requests.PreparedRequest) -> Optional[str]:
633        """Recover the token a request was signed with from its auth header.
634
635        The prefix is checked rather than assumed: the header may have been written by
636        something other than this authenticator, and slicing blindly would leave membership in
637        `_states` as the only thing standing between a mangled value and a wrong attribution.
638        """
639        value = request.headers.get(self._header)
640        if not value:
641            return None
642        if self._auth_method:
643            prefix = f"{self._auth_method} "
644            if not value.startswith(prefix):
645                return None
646            token = value[len(prefix) :].strip()
647        else:
648            token = value.strip()
649        return token if token in self._states else None
650
651    @staticmethod
652    def _header_int(response: requests.Response, header: Optional[str]) -> Optional[int]:
653        if not header:
654            return None
655        value = response.headers.get(header)
656        if value is None:
657            return None
658        try:
659            return int(value)
660        except (TypeError, ValueError):
661            return None
662
663    @staticmethod
664    def _header_datetime(
665        response: requests.Response, header: Optional[str]
666    ) -> Optional[AirbyteDateTime]:
667        if not header:
668            return None
669        value = response.headers.get(header)
670        if value is None:
671            return None
672        try:
673            # Same parsing rules as `reset_path`, so epoch seconds and ISO 8601 both work.
674            return ab_datetime_parse(value)
675        except Exception:
676            return None