airbyte_cdk.sources.streams.http.requests_native_auth.abstract_oauth

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4
  5import logging
  6import threading
  7from abc import abstractmethod
  8from datetime import timedelta
  9from json import JSONDecodeError
 10from typing import Any, List, Mapping, MutableMapping, Optional, Tuple, Union
 11
 12import backoff
 13import requests
 14from requests.auth import AuthBase
 15
 16from airbyte_cdk.models import FailureType, Level
 17from airbyte_cdk.sources.http_logger import format_http_message
 18from airbyte_cdk.sources.message import MessageRepository, NoopMessageRepository
 19from airbyte_cdk.utils import AirbyteTracedException
 20from airbyte_cdk.utils.airbyte_secrets_utils import add_to_secrets
 21from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now, ab_datetime_parse
 22
 23from ..exceptions import DefaultBackoffException
 24
 25logger = logging.getLogger("airbyte")
 26_NOOP_MESSAGE_REPOSITORY = NoopMessageRepository()
 27
 28
 29class ResponseKeysMaxRecurtionReached(AirbyteTracedException):
 30    """
 31    Raised when the max level of recursion is reached, when trying to
 32    find-and-get the target key, during the `_make_handled_request`
 33    """
 34
 35
 36class AbstractOauth2Authenticator(AuthBase):
 37    """
 38    Abstract class for an OAuth authenticators that implements the OAuth token refresh flow. The authenticator
 39    is designed to generically perform the refresh flow without regard to how config fields are get/set by
 40    delegating that behavior to the classes implementing the interface.
 41    """
 42
 43    _NO_STREAM_NAME = None
 44
 45    # Class-level lock to prevent concurrent token refresh across multiple authenticator instances.
 46    # This is necessary because multiple streams may share the same OAuth credentials (refresh token)
 47    # through the connector config. Without this lock, concurrent refresh attempts can cause race
 48    # conditions where one stream successfully refreshes the token while others fail because the
 49    # refresh token has been invalidated (especially for single-use refresh tokens).
 50    _token_refresh_lock: threading.Lock = threading.Lock()
 51
 52    def __init__(
 53        self,
 54        refresh_token_error_status_codes: Tuple[int, ...] = (),
 55        refresh_token_error_key: str = "",
 56        refresh_token_error_values: Tuple[str, ...] = (),
 57    ) -> None:
 58        """
 59        If all of refresh_token_error_status_codes, refresh_token_error_key, and refresh_token_error_values are set,
 60        then http errors with such params will be wrapped in AirbyteTracedException.
 61        """
 62        self._refresh_token_error_status_codes = refresh_token_error_status_codes
 63        self._refresh_token_error_key = refresh_token_error_key
 64        self._refresh_token_error_values = refresh_token_error_values
 65
 66    def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest:
 67        """Attach the HTTP headers required to authenticate on the HTTP request"""
 68        request.headers.update(self.get_auth_header())
 69        return request
 70
 71    @property
 72    def _is_access_token_flow(self) -> bool:
 73        return self.get_token_refresh_endpoint() is None and self.access_token is not None
 74
 75    @property
 76    def token_expiry_is_time_of_expiration(self) -> bool:
 77        """
 78        Indicates that the Token Expiry returns the date until which the token will be valid, not the amount of time it will be valid.
 79        """
 80
 81        return False
 82
 83    @property
 84    def token_expiry_date_format(self) -> Optional[str]:
 85        """
 86        Format of the datetime; exists it if expires_in is returned as the expiration datetime instead of seconds until it expires
 87        """
 88
 89        return None
 90
 91    def get_auth_header(self) -> Mapping[str, Any]:
 92        """HTTP header to set on the requests"""
 93        token = self.access_token if self._is_access_token_flow else self.get_access_token()
 94        return {"Authorization": f"Bearer {token}"}
 95
 96    def get_access_token(self) -> str:
 97        """
 98        Returns the access token.
 99
100        This method uses double-checked locking to ensure thread-safe token refresh.
101        When multiple threads (streams) detect an expired token simultaneously, only one
102        will perform the refresh while others wait. After acquiring the lock, the token
103        expiry is re-checked to avoid redundant refresh attempts.
104        """
105        if self.token_has_expired():
106            with self._token_refresh_lock:
107                # Double-check after acquiring lock - another thread may have already refreshed
108                if self.token_has_expired():
109                    self.refresh_and_set_access_token()
110
111        return self.access_token
112
113    def refresh_and_set_access_token(self) -> None:
114        """Force refresh the access token and update internal state.
115
116        This method refreshes the access token regardless of whether it has expired,
117        and updates the internal token and expiry date. Subclasses may override this
118        to handle additional state updates (e.g., persisting new refresh tokens).
119        """
120        token, expires_in = self.refresh_access_token()
121        self.access_token = token
122        self.set_token_expiry_date(expires_in)
123
124    def token_has_expired(self) -> bool:
125        """Returns True if the token is expired"""
126        return ab_datetime_now() > self.get_token_expiry_date()
127
128    def _build_standard_refresh_args(self) -> MutableMapping[str, Any]:
129        """Build the standard OAuth refresh args (grant_type, refresh_token, client
130        credentials, scopes, plus any user-configured `refresh_request_body` extras).
131
132        Used by both `build_refresh_request_body()` and
133        `build_refresh_request_query_params()` so the same set of args can be emitted
134        in either the body or the URL query string depending on
135        `should_send_refresh_request_as_query_params()`.
136
137        Client credentials (client_id and client_secret) are excluded when
138        `refresh_request_headers` contains an `Authorization` header (e.g. Basic
139        auth). This is required by OAuth providers like Gong that expect credentials
140        ONLY in the Authorization header and reject requests that include them in
141        both places.
142        """
143        headers = self.get_refresh_request_headers()
144        credentials_in_header = headers and "Authorization" in headers
145        include_client_credentials = not credentials_in_header
146
147        payload: MutableMapping[str, Any] = {
148            self.get_grant_type_name(): self.get_grant_type(),
149        }
150
151        if include_client_credentials:
152            payload[self.get_client_id_name()] = self.get_client_id()
153            payload[self.get_client_secret_name()] = self.get_client_secret()
154
155        payload[self.get_refresh_token_name()] = self.get_refresh_token()
156
157        if self.get_scopes():
158            payload["scopes"] = self.get_scopes()
159
160        if self.get_refresh_request_body():
161            for key, val in self.get_refresh_request_body().items():
162                # Existing oauth args take precedence over custom configured fields.
163                if key not in payload:
164                    payload[key] = val
165
166        return payload
167
168    def build_refresh_request_body(self) -> Mapping[str, Any]:
169        """Returns the request body to set on the refresh request.
170
171        When `should_send_refresh_request_as_query_params()` is `True`, the standard
172        refresh args are emitted on the URL query string instead and this method
173        returns an empty body. This supports OAuth providers like Gong that document
174        their refresh endpoint as a `POST` with parameters on the URL query string
175        and an empty body.
176        """
177        if self.should_send_refresh_request_as_query_params():
178            return {}
179        return self._build_standard_refresh_args()
180
181    def build_refresh_request_headers(self) -> Mapping[str, Any] | None:
182        """
183        Returns the request headers to set on the refresh request
184
185        """
186        headers = self.get_refresh_request_headers()
187        return headers if headers else None
188
189    def build_refresh_request_query_params(self) -> Mapping[str, Any] | None:
190        """Returns the URL query string parameters to set on the refresh request.
191
192        When `should_send_refresh_request_as_query_params()` is `True`, the standard
193        refresh args (grant_type, refresh_token, client credentials, scopes, plus
194        any user-configured `refresh_request_body` extras) are returned here and
195        `build_refresh_request_body()` returns an empty body.
196
197        Returns `None` otherwise so existing authenticators retain their previous
198        behavior (no query params on the refresh URL).
199        """
200        if not self.should_send_refresh_request_as_query_params():
201            return None
202        return self._build_standard_refresh_args()
203
204    def refresh_access_token(self) -> Tuple[str, AirbyteDateTime]:
205        """
206        Returns the refresh token and its expiration datetime
207
208        :return: a tuple of (access_token, token_lifespan)
209        """
210        try:
211            response_json = self._make_handled_request()
212        except (
213            requests.exceptions.ConnectionError,
214            requests.exceptions.ConnectTimeout,
215            requests.exceptions.ReadTimeout,
216        ) as e:
217            raise AirbyteTracedException(
218                message="OAuth access token refresh request failed due to a network error.",
219                internal_message=f"Network error during OAuth token refresh after retries were exhausted: {e}",
220                failure_type=FailureType.transient_error,
221            ) from e
222        self._ensure_access_token_in_response(response_json)
223
224        return (
225            self._extract_access_token(response_json),
226            self._extract_token_expiry_date(response_json),
227        )
228
229    # ----------------
230    # PRIVATE METHODS
231    # ----------------
232
233    def _default_token_expiry_date(self) -> AirbyteDateTime:
234        """
235        Returns the default token expiry date
236        """
237        # 1 hour was chosen as a middle ground to avoid unnecessary frequent refreshes and token expiration
238        default_token_expiry_duration_hours = 1  # 1 hour
239        return ab_datetime_now() + timedelta(hours=default_token_expiry_duration_hours)
240
241    def _wrap_refresh_token_exception(
242        self, exception: requests.exceptions.RequestException
243    ) -> bool:
244        """
245        Wraps and handles exceptions that occur during the refresh token process.
246
247        This method checks if the provided exception is related to a refresh token error
248        by examining the response status code and specific error content.
249
250        Args:
251            exception (requests.exceptions.RequestException): The exception raised during the request.
252
253        Returns:
254            bool: True if the exception is related to a refresh token error, False otherwise.
255        """
256        try:
257            if exception.response is not None:
258                exception_content = exception.response.json()
259            else:
260                return False
261        except JSONDecodeError:
262            return False
263        return (
264            exception.response.status_code in self._refresh_token_error_status_codes
265            and exception_content.get(self._refresh_token_error_key)
266            in self._refresh_token_error_values
267        )
268
269    @backoff.on_exception(
270        backoff.expo,
271        (
272            DefaultBackoffException,
273            requests.exceptions.ConnectionError,
274            requests.exceptions.ConnectTimeout,
275            requests.exceptions.ReadTimeout,
276        ),
277        on_backoff=lambda details: logger.info(
278            f"Caught retryable error after {details['tries']} tries. Waiting {details['wait']} seconds then retrying..."
279        ),
280        max_time=300,
281    )
282    def _make_handled_request(self) -> Any:
283        """
284        Makes a handled HTTP request to refresh an OAuth token.
285
286        This method sends a POST request to the token refresh endpoint with the necessary
287        headers and body to obtain a new access token. It handles various exceptions that
288        may occur during the request and logs the response for troubleshooting purposes.
289
290        Returns:
291            Mapping[str, Any]: The JSON response from the token refresh endpoint.
292
293        Raises:
294            DefaultBackoffException: If the response status code is 429 (Too Many Requests)
295                                     or any 5xx server error.
296            AirbyteTracedException: If the refresh token is invalid or expired, prompting
297                                    re-authentication.
298            Exception: For any other exceptions that occur during the request.
299        """
300        try:
301            response = requests.request(
302                method="POST",
303                url=self.get_token_refresh_endpoint(),  # type: ignore # returns None, if not provided, but str | bytes is expected.
304                data=self.build_refresh_request_body(),
305                headers=self.build_refresh_request_headers(),
306                params=self.build_refresh_request_query_params(),
307            )
308
309            if not response.ok:
310                # log the response even if the request failed for troubleshooting purposes
311                self._log_response(response)
312                response.raise_for_status()
313
314            response_json = response.json()
315
316            try:
317                # extract the access token and add to secrets to avoid logging the raw value
318                access_key = self._extract_access_token(response_json)
319                if access_key:
320                    add_to_secrets(access_key)
321            except ResponseKeysMaxRecurtionReached as e:
322                # could not find the access token in the response, so do nothing
323                pass
324
325            self._log_response(response)
326
327            return response_json
328        except requests.exceptions.RequestException as e:
329            if e.response is not None:
330                if e.response.status_code == 429 or e.response.status_code >= 500:
331                    raise DefaultBackoffException(
332                        request=e.response.request,
333                        response=e.response,
334                        failure_type=FailureType.transient_error,
335                    )
336            if self._wrap_refresh_token_exception(e):
337                response_info = (
338                    f"HTTP {e.response.status_code}: {e.response.text[:1000]}"
339                    if e.response is not None
340                    else str(e)
341                )
342                raise AirbyteTracedException(
343                    internal_message=f"Refresh token rejected by the OAuth token endpoint. {response_info}",
344                    message="Refresh token was rejected by the OAuth provider (invalid, expired, or already used). Re-authenticate this source's credentials in its connection settings.",
345                    failure_type=FailureType.config_error,
346                ) from e
347            raise
348        except Exception as e:
349            raise AirbyteTracedException(
350                message="OAuth access token refresh request failed.",
351                internal_message=f"Unexpected error during OAuth token refresh: {e}",
352                failure_type=FailureType.system_error,
353            ) from e
354
355    def _ensure_access_token_in_response(self, response_data: Mapping[str, Any]) -> None:
356        """
357        Ensures that the access token is present in the response data.
358
359        This method attempts to extract the access token from the provided response data.
360        If the access token is not found, it raises an exception indicating that the token
361        refresh API response was missing the access token.
362
363        Args:
364            response_data (Mapping[str, Any]): The response data from which to extract the access token.
365
366        Raises:
367            Exception: If the access token is not found in the response data.
368            ResponseKeysMaxRecurtionReached: If the maximum recursion depth is reached while extracting the access token.
369        """
370        try:
371            access_key = self._extract_access_token(response_data)
372            if not access_key:
373                raise Exception(
374                    f"Token refresh API response was missing access token {self.get_access_token_name()}"
375                )
376        except ResponseKeysMaxRecurtionReached as e:
377            raise e
378
379    def _parse_token_expiration_date(self, value: Union[str, int]) -> AirbyteDateTime:
380        """
381        Parse a string or integer token expiration date into a datetime object
382
383        :return: expiration datetime
384        """
385        if self.token_expiry_is_time_of_expiration:
386            if not self.token_expiry_date_format:
387                raise ValueError(
388                    f"Invalid token expiry date format {self.token_expiry_date_format}; a string representing the format is required."
389                )
390            try:
391                return ab_datetime_parse(str(value))
392            except ValueError as e:
393                raise ValueError(f"Invalid token expiry date format: {e}")
394        else:
395            try:
396                # Only accept numeric values (as int/float/string) when no format specified
397                seconds = int(float(str(value)))
398                return ab_datetime_now() + timedelta(seconds=seconds)
399            except (ValueError, TypeError):
400                raise ValueError(
401                    f"Invalid expires_in value: {value}. Expected number of seconds when no format specified."
402                )
403
404    def _extract_access_token(self, response_data: Mapping[str, Any]) -> Any:
405        """
406        Extracts the access token from the given response data.
407
408        Args:
409            response_data (Mapping[str, Any]): The response data from which to extract the access token.
410
411        Returns:
412            str: The extracted access token.
413        """
414        return self._find_and_get_value_from_response(response_data, self.get_access_token_name())
415
416    def _extract_refresh_token(self, response_data: Mapping[str, Any]) -> Any:
417        """
418        Extracts the refresh token from the given response data.
419
420        Args:
421            response_data (Mapping[str, Any]): The response data from which to extract the refresh token.
422
423        Returns:
424            str: The extracted refresh token.
425        """
426        return self._find_and_get_value_from_response(response_data, self.get_refresh_token_name())
427
428    def _extract_token_expiry_date(self, response_data: Mapping[str, Any]) -> AirbyteDateTime:
429        """
430        Extracts the token_expiry_date, like `expires_in` or `expires_at`, etc from the given response data.
431
432        If the token_expiry_date is not found, it will return an existing token expiry date if set, or a default token expiry date.
433
434        Args:
435            response_data (Mapping[str, Any]): The response data from which to extract the token_expiry_date.
436
437        Returns:
438            The extracted token_expiry_date or None if not found.
439        """
440        expires_in = self._find_and_get_value_from_response(
441            response_data, self.get_expires_in_name()
442        )
443        if expires_in is not None:
444            return self._parse_token_expiration_date(expires_in)
445
446        # expires_in is None
447        existing_expiry_date = self.get_token_expiry_date()
448        if existing_expiry_date and not self.token_has_expired():
449            return existing_expiry_date
450
451        return self._default_token_expiry_date()
452
453    def _find_and_get_value_from_response(
454        self,
455        response_data: Mapping[str, Any],
456        key_name: str,
457        max_depth: int = 5,
458        current_depth: int = 0,
459    ) -> Any:
460        """
461        Recursively searches for a specified key in a nested dictionary or list and returns its value if found.
462
463        Args:
464            response_data (Mapping[str, Any]): The response data to search through, which can be a dictionary or a list.
465            key_name (str): The key to search for in the response data.
466            max_depth (int, optional): The maximum depth to search for the key to avoid infinite recursion. Defaults to 5.
467            current_depth (int, optional): The current depth of the recursion. Defaults to 0.
468
469        Returns:
470            Any: The value associated with the specified key if found, otherwise None.
471
472        Raises:
473            AirbyteTracedException: If the maximum recursion depth is reached without finding the key.
474        """
475        if current_depth > max_depth:
476            # this is needed to avoid an inf loop, possible with a very deep nesting observed.
477            message = f"The maximum level of recursion is reached. Couldn't find the specified `{key_name}` in the response."
478            raise ResponseKeysMaxRecurtionReached(
479                internal_message=message, message=message, failure_type=FailureType.config_error
480            )
481
482        if isinstance(response_data, dict):
483            # get from the root level
484            if key_name in response_data:
485                return response_data[key_name]
486
487            # get from the nested object
488            for _, value in response_data.items():
489                result = self._find_and_get_value_from_response(
490                    value, key_name, max_depth, current_depth + 1
491                )
492                if result is not None:
493                    return result
494
495        # get from the nested array object
496        elif isinstance(response_data, list):
497            for item in response_data:
498                result = self._find_and_get_value_from_response(
499                    item, key_name, max_depth, current_depth + 1
500                )
501                if result is not None:
502                    return result
503
504        return None
505
506    @property
507    def _message_repository(self) -> Optional[MessageRepository]:
508        """
509        The implementation can define a message_repository if it wants debugging logs for HTTP requests
510        """
511        return _NOOP_MESSAGE_REPOSITORY
512
513    def _log_response(self, response: requests.Response) -> None:
514        """
515        Logs the HTTP response using the message repository if it is available.
516
517        Args:
518            response (requests.Response): The HTTP response to log.
519        """
520        if self._message_repository:
521            self._message_repository.log_message(
522                Level.DEBUG,
523                lambda: format_http_message(
524                    response,
525                    "Refresh token",
526                    "Obtains access token",
527                    self._NO_STREAM_NAME,
528                    is_auxiliary=True,
529                    type="AUTH",
530                ),
531            )
532
533    # ----------------
534    # ABSTR METHODS
535    # ----------------
536
537    @abstractmethod
538    def get_token_refresh_endpoint(self) -> Optional[str]:
539        """Returns the endpoint to refresh the access token"""
540
541    @abstractmethod
542    def get_client_id_name(self) -> str:
543        """The client id name to authenticate"""
544
545    @abstractmethod
546    def get_client_id(self) -> str:
547        """The client id to authenticate"""
548
549    @abstractmethod
550    def get_client_secret_name(self) -> str:
551        """The client secret name to authenticate"""
552
553    @abstractmethod
554    def get_client_secret(self) -> str:
555        """The client secret to authenticate"""
556
557    @abstractmethod
558    def get_refresh_token_name(self) -> str:
559        """The refresh token name to authenticate"""
560
561    @abstractmethod
562    def get_refresh_token(self) -> Optional[str]:
563        """The token used to refresh the access token when it expires"""
564
565    @abstractmethod
566    def get_scopes(self) -> List[str]:
567        """List of requested scopes"""
568
569    @abstractmethod
570    def get_token_expiry_date(self) -> AirbyteDateTime:
571        """Expiration date of the access token"""
572
573    @abstractmethod
574    def set_token_expiry_date(self, value: AirbyteDateTime) -> None:
575        """Setter for access token expiration date"""
576
577    @abstractmethod
578    def get_access_token_name(self) -> str:
579        """Field to extract access token from in the response"""
580
581    @abstractmethod
582    def get_expires_in_name(self) -> str:
583        """Returns the expires_in field name"""
584
585    @abstractmethod
586    def get_refresh_request_body(self) -> Mapping[str, Any]:
587        """Returns the request body to set on the refresh request"""
588
589    @abstractmethod
590    def get_refresh_request_headers(self) -> Mapping[str, Any]:
591        """Returns the request headers to set on the refresh request"""
592
593    def should_send_refresh_request_as_query_params(self) -> bool:
594        """Returns `True` if the standard refresh args should be sent on the URL
595        query string instead of in the request body.
596
597        Defaults to `False` so existing authenticators retain their previous
598        behavior (params in body, no query params on the refresh URL). Subclasses
599        can override this to opt into the URL-query-string shape required by OAuth
600        providers like Gong.
601        """
602        return False
603
604    @abstractmethod
605    def get_grant_type(self) -> str:
606        """Returns grant_type specified for requesting access_token"""
607
608    @abstractmethod
609    def get_grant_type_name(self) -> str:
610        """Returns grant_type specified name for requesting access_token"""
611
612    @property
613    @abstractmethod
614    def access_token(self) -> str:
615        """Returns the access token"""
616
617    @access_token.setter
618    @abstractmethod
619    def access_token(self, value: str) -> str:
620        """Setter for the access token"""