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

 48class RateLimitedMultipleTokenAuthenticator(DeclarativeAuthenticator):
 49    """Authenticator that rotates between multiple interchangeable tokens with per-token quota tracking.
 50
 51    Each outgoing request is classified into a quota pool using the pool's request matchers.
 52    The active token's counter for the matched pool is decremented locally; when it is exhausted
 53    the authenticator rotates to the next token. When all tokens are exhausted for a pool, it
 54    waits until the earliest quota reset (bounded by `max_wait_time`) and then refreshes all
 55    counters from `quota_status_url`, or raises a transient error if the wait would be too long.
 56
 57    A proactive throttling budget spreads the last calls over the time remaining until reset:
 58    once every token's remaining count for a pool drops below its reserve
 59    (`max(budget_min_reserve, budget_reserve_fraction * limit)`), a small delay proportional to
 60    `seconds_until_reset / total_remaining` (capped at 10s) is injected before each request.
 61
 62    Counters are seeded per token from `quota_status_url` on first use and refreshed after an
 63    exhaustion wait. All state transitions are guarded by a lock so the authenticator can be
 64    shared safely across concurrent streams; sleeps never hold the lock.
 65    """
 66
 67    HEARTBEAT_INTERVAL = 60.0  # Log every 60s during exhaustion wait
 68    MAX_BUDGET_DELAY = 10.0  # Cap for the per-request proactive throttling delay
 69    MIN_EXHAUSTION_WAIT = 5.0  # Floor for the exhaustion wait, so stale reset timestamps can't cause a refresh busy-loop
 70
 71    def __init__(
 72        self,
 73        tokens: List[str],
 74        quotas: List[TokenQuota],
 75        quota_status_url: str,
 76        quota_status_http_method: str = "GET",
 77        quota_status_headers: Optional[Mapping[str, str]] = None,
 78        auth_method: str = "Bearer",
 79        header: str = "Authorization",
 80        max_wait_time: timedelta = timedelta(hours=2),
 81        budget_reserve_fraction: float = 0.1,
 82        budget_min_reserve: int = 50,
 83    ) -> None:
 84        if not tokens:
 85            raise AirbyteTracedException(
 86                failure_type=FailureType.config_error,
 87                internal_message="RateLimitedMultipleTokenAuthenticator requires at least one token",
 88                message="Authentication tokens are missing from the configuration.",
 89            )
 90        if not quotas:
 91            raise AirbyteTracedException(
 92                failure_type=FailureType.config_error,
 93                internal_message="RateLimitedMultipleTokenAuthenticator requires at least one quota pool",
 94                message="Quota pool configuration is missing.",
 95            )
 96        self._logger = logging.getLogger("airbyte")
 97        self._tokens = list(tokens)
 98        self._quotas = quotas
 99        self._quota_status_url = quota_status_url
100        self._quota_status_http_method = quota_status_http_method
101        self._quota_status_headers = dict(quota_status_headers or {})
102        self._auth_method = auth_method
103        self._header = header
104        self._max_wait_time = max_wait_time
105        self._budget_reserve_fraction = budget_reserve_fraction
106        self._budget_min_reserve = budget_min_reserve
107
108        self._lock = threading.RLock()
109        self._refresh_lock = threading.Lock()
110        self._initialized = False
111        self._budget_logged = False
112        self._unmatched_logged = False
113        self._states: dict[str, dict[str, _QuotaState]] = {}
114        self._token_to_http_client: Mapping[str, HttpClient] = {
115            token: HttpClient(
116                name="quota_status",
117                logger=self._logger,
118                authenticator=TokenAuthenticator(
119                    token, auth_method=self._auth_method, auth_header=self._header
120                ),
121                use_cache=False,  # quota values change frequently; never reuse cached responses
122            )
123            for token in self._tokens
124        }
125        self._tokens_iter = cycle(self._tokens)
126        self._active_token = next(self._tokens_iter)
127
128    @property
129    def auth_header(self) -> str:
130        return self._header
131
132    @property
133    def token(self) -> str:
134        with self._lock:
135            return f"{self._auth_method} {self._active_token}".strip()
136
137    def __call__(self, request: requests.PreparedRequest) -> Any:
138        """Attach the HTTP headers required to authenticate on the HTTP request"""
139        self._ensure_initialized()
140        quota = self._match_quota(request)
141        token = self._acquire_call(quota)
142        request.headers[self._header] = f"{self._auth_method} {token}".strip()
143        return request
144
145    def _ensure_initialized(self) -> None:
146        if self._initialized:
147            return
148        with self._refresh_lock:
149            if not self._initialized:
150                self._seed_all_tokens()
151                self._initialized = True
152
153    def _match_quota(self, request: requests.PreparedRequest) -> TokenQuota:
154        default_quota: Optional[TokenQuota] = None
155        for quota in self._quotas:
156            if quota.matchers:
157                if any(matcher(request) for matcher in quota.matchers):
158                    return quota
159            elif default_quota is None:
160                default_quota = quota
161        if default_quota is None:
162            if not self._unmatched_logged:
163                self._logger.warning(
164                    "Request %s did not match any quota pool; falling back to '%s'. Consider defining a matcher-less default pool.",
165                    request.url,
166                    self._quotas[0].name,
167                )
168                self._unmatched_logged = True
169        return default_quota or self._quotas[0]
170
171    def _acquire_call(self, quota: TokenQuota) -> str:
172        """Reserve one call from the matched pool and return the token it was charged to.
173
174        `max_wait_time` bounds the *total* time spent waiting across all refresh attempts of a
175        single exhaustion episode, so stale reset timestamps cannot cause an endless reseed loop.
176        """
177        exhaustion_deadline: Optional[float] = None
178        while True:
179            budget_delay: Optional[float] = None
180            wait_for_reset: Optional[float] = None
181            with self._lock:
182                token = self._active_token
183                state = self._states[token][quota.name]
184                if state.remaining > 0:
185                    state.remaining -= 1
186                    budget_delay = self._compute_budget_delay(quota)
187                elif all(self._states[token][quota.name].remaining <= 0 for token in self._tokens):
188                    now = time.monotonic()
189                    if exhaustion_deadline is None:
190                        exhaustion_deadline = now + self._max_wait_time.total_seconds()
191                    remaining_budget = exhaustion_deadline - now
192                    min_time_to_wait = min(
193                        (
194                            self._states[token][quota.name].reset_at - ab_datetime_now()
195                        ).total_seconds()
196                        for token in self._tokens
197                    )
198                    if remaining_budget <= 0 or min_time_to_wait >= remaining_budget:
199                        raise AirbyteTracedException(
200                            failure_type=FailureType.transient_error,
201                            internal_message=f"Rate limits for all tokens (quota: {quota.name}) were reached and the next reset exceeds max_wait_time",
202                            message="Rate limit is exceeded for all provided tokens.",
203                        )
204                    wait_for_reset = min(
205                        max(min_time_to_wait, self.MIN_EXHAUSTION_WAIT),
206                        remaining_budget,
207                    )
208                else:
209                    self._active_token = next(self._tokens_iter)
210                    continue
211
212            if wait_for_reset is not None:
213                self._logger.info(
214                    "All tokens exhausted (quota: %s). Waiting %.0fs until rate limit resets.",
215                    quota.name,
216                    wait_for_reset,
217                )
218                self._sleep_with_heartbeat(wait_for_reset, quota.name)
219                self._refresh_after_exhaustion(quota)
220                continue
221
222            if budget_delay is not None and budget_delay >= 0.1:
223                if not self._budget_logged:
224                    self._logger.info(
225                        "API budget: throttling requests (%.1fs delay) for quota '%s'.",
226                        budget_delay,
227                        quota.name,
228                    )
229                    self._budget_logged = True
230                time.sleep(budget_delay)
231            return token
232
233    def _compute_budget_delay(self, quota: TokenQuota) -> Optional[float]:
234        """Compute the proactive throttling delay. Must be called while holding the lock."""
235        states = [self._states[token][quota.name] for token in self._tokens]
236        if not all(state.remaining <= self._get_budget_reserve(state) for state in states):
237            return None
238
239        active_state = self._states[self._active_token][quota.name]
240        seconds_to_reset = max((active_state.reset_at - ab_datetime_now()).total_seconds(), 0)
241        total_remaining = sum(max(state.remaining, 0) for state in states)
242        if total_remaining <= 0 or seconds_to_reset <= 0:
243            return None
244
245        return min(seconds_to_reset / total_remaining, self.MAX_BUDGET_DELAY)
246
247    def _get_budget_reserve(self, state: _QuotaState) -> float:
248        return max(self._budget_min_reserve, state.limit * self._budget_reserve_fraction)
249
250    def _sleep_with_heartbeat(self, total_seconds: float, quota_name: str) -> None:
251        """Sleep for `total_seconds` in chunks, logging progress so operators can see the connector is not stuck."""
252        remaining = total_seconds
253        while remaining > 0:
254            chunk = min(remaining, self.HEARTBEAT_INTERVAL)
255            time.sleep(chunk)
256            remaining -= chunk
257            if remaining > 0:
258                self._logger.info(
259                    "Rate limit exhausted (quota: %s). Waiting for reset — %.0fs remaining.",
260                    quota_name,
261                    remaining,
262                )
263
264    def _refresh_after_exhaustion(self, quota: TokenQuota) -> None:
265        """Refresh counters after an exhaustion wait. Only one thread refreshes; others re-check state."""
266        with self._refresh_lock:
267            with self._lock:
268                still_exhausted = all(
269                    self._states[token][quota.name].remaining <= 0 for token in self._tokens
270                )
271            if still_exhausted:
272                self._seed_all_tokens()
273
274    def _seed_all_tokens(self) -> None:
275        # The wholesale swap intentionally discards local decrements made by concurrent threads
276        # between the fetch and the swap: the server response is the closest thing to truth, and
277        # merging local decrements on top of it would double-count the calls the server has
278        # already observed. The worst case is a slight overcount of `remaining` (requests that
279        # were in flight during the fetch), which at most causes a few 429s near the quota
280        # boundary that the stream-level error handler absorbs.
281        states = {token: self._fetch_quota_states(token) for token in self._tokens}
282        with self._lock:
283            self._states = states
284            self._budget_logged = False
285
286    def _fetch_quota_states(self, token: str) -> dict[str, _QuotaState]:
287        http_client = self._token_to_http_client[token]
288        _, response = http_client.send_request(
289            http_method=self._quota_status_http_method,
290            url=self._quota_status_url,
291            headers=self._quota_status_headers,
292            request_kwargs={},
293        )
294        response_body = response.json()
295
296        states = {}
297        for quota in self._quotas:
298            remaining = self._extract_path(response_body, quota.remaining_path)
299            reset = self._extract_path(response_body, quota.reset_path)
300            limit = (
301                self._extract_path(response_body, quota.limit_path)
302                if quota.limit_path
303                else remaining
304            )
305            states[quota.name] = _QuotaState(
306                remaining=int(remaining),
307                reset_at=ab_datetime_parse(reset),
308                limit=int(limit),
309            )
310        return states
311
312    def _extract_path(self, response_body: Mapping[str, Any], path: List[str]) -> Any:
313        value: Any = response_body
314        for key in path:
315            if not isinstance(value, Mapping) or key not in value:
316                raise AirbyteTracedException(
317                    failure_type=FailureType.config_error,
318                    internal_message=f"Quota status response did not contain expected path: {path}",
319                    message="Quota status response is missing an expected field.",
320                )
321            value = value[key]
322        return value