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                message = "Refresh token is invalid or expired. Please re-authenticate from Sources/<your source>/Settings."
338                raise AirbyteTracedException(
339                    internal_message=message, message=message, failure_type=FailureType.config_error
340                )
341            raise
342        except Exception as e:
343            raise AirbyteTracedException(
344                message="OAuth access token refresh request failed.",
345                internal_message=f"Unexpected error during OAuth token refresh: {e}",
346                failure_type=FailureType.system_error,
347            ) from e
348
349    def _ensure_access_token_in_response(self, response_data: Mapping[str, Any]) -> None:
350        """
351        Ensures that the access token is present in the response data.
352
353        This method attempts to extract the access token from the provided response data.
354        If the access token is not found, it raises an exception indicating that the token
355        refresh API response was missing the access token.
356
357        Args:
358            response_data (Mapping[str, Any]): The response data from which to extract the access token.
359
360        Raises:
361            Exception: If the access token is not found in the response data.
362            ResponseKeysMaxRecurtionReached: If the maximum recursion depth is reached while extracting the access token.
363        """
364        try:
365            access_key = self._extract_access_token(response_data)
366            if not access_key:
367                raise Exception(
368                    f"Token refresh API response was missing access token {self.get_access_token_name()}"
369                )
370        except ResponseKeysMaxRecurtionReached as e:
371            raise e
372
373    def _parse_token_expiration_date(self, value: Union[str, int]) -> AirbyteDateTime:
374        """
375        Parse a string or integer token expiration date into a datetime object
376
377        :return: expiration datetime
378        """
379        if self.token_expiry_is_time_of_expiration:
380            if not self.token_expiry_date_format:
381                raise ValueError(
382                    f"Invalid token expiry date format {self.token_expiry_date_format}; a string representing the format is required."
383                )
384            try:
385                return ab_datetime_parse(str(value))
386            except ValueError as e:
387                raise ValueError(f"Invalid token expiry date format: {e}")
388        else:
389            try:
390                # Only accept numeric values (as int/float/string) when no format specified
391                seconds = int(float(str(value)))
392                return ab_datetime_now() + timedelta(seconds=seconds)
393            except (ValueError, TypeError):
394                raise ValueError(
395                    f"Invalid expires_in value: {value}. Expected number of seconds when no format specified."
396                )
397
398    def _extract_access_token(self, response_data: Mapping[str, Any]) -> Any:
399        """
400        Extracts the access token from the given response data.
401
402        Args:
403            response_data (Mapping[str, Any]): The response data from which to extract the access token.
404
405        Returns:
406            str: The extracted access token.
407        """
408        return self._find_and_get_value_from_response(response_data, self.get_access_token_name())
409
410    def _extract_refresh_token(self, response_data: Mapping[str, Any]) -> Any:
411        """
412        Extracts the refresh token from the given response data.
413
414        Args:
415            response_data (Mapping[str, Any]): The response data from which to extract the refresh token.
416
417        Returns:
418            str: The extracted refresh token.
419        """
420        return self._find_and_get_value_from_response(response_data, self.get_refresh_token_name())
421
422    def _extract_token_expiry_date(self, response_data: Mapping[str, Any]) -> AirbyteDateTime:
423        """
424        Extracts the token_expiry_date, like `expires_in` or `expires_at`, etc from the given response data.
425
426        If the token_expiry_date is not found, it will return an existing token expiry date if set, or a default token expiry date.
427
428        Args:
429            response_data (Mapping[str, Any]): The response data from which to extract the token_expiry_date.
430
431        Returns:
432            The extracted token_expiry_date or None if not found.
433        """
434        expires_in = self._find_and_get_value_from_response(
435            response_data, self.get_expires_in_name()
436        )
437        if expires_in is not None:
438            return self._parse_token_expiration_date(expires_in)
439
440        # expires_in is None
441        existing_expiry_date = self.get_token_expiry_date()
442        if existing_expiry_date and not self.token_has_expired():
443            return existing_expiry_date
444
445        return self._default_token_expiry_date()
446
447    def _find_and_get_value_from_response(
448        self,
449        response_data: Mapping[str, Any],
450        key_name: str,
451        max_depth: int = 5,
452        current_depth: int = 0,
453    ) -> Any:
454        """
455        Recursively searches for a specified key in a nested dictionary or list and returns its value if found.
456
457        Args:
458            response_data (Mapping[str, Any]): The response data to search through, which can be a dictionary or a list.
459            key_name (str): The key to search for in the response data.
460            max_depth (int, optional): The maximum depth to search for the key to avoid infinite recursion. Defaults to 5.
461            current_depth (int, optional): The current depth of the recursion. Defaults to 0.
462
463        Returns:
464            Any: The value associated with the specified key if found, otherwise None.
465
466        Raises:
467            AirbyteTracedException: If the maximum recursion depth is reached without finding the key.
468        """
469        if current_depth > max_depth:
470            # this is needed to avoid an inf loop, possible with a very deep nesting observed.
471            message = f"The maximum level of recursion is reached. Couldn't find the specified `{key_name}` in the response."
472            raise ResponseKeysMaxRecurtionReached(
473                internal_message=message, message=message, failure_type=FailureType.config_error
474            )
475
476        if isinstance(response_data, dict):
477            # get from the root level
478            if key_name in response_data:
479                return response_data[key_name]
480
481            # get from the nested object
482            for _, value in response_data.items():
483                result = self._find_and_get_value_from_response(
484                    value, key_name, max_depth, current_depth + 1
485                )
486                if result is not None:
487                    return result
488
489        # get from the nested array object
490        elif isinstance(response_data, list):
491            for item in response_data:
492                result = self._find_and_get_value_from_response(
493                    item, key_name, max_depth, current_depth + 1
494                )
495                if result is not None:
496                    return result
497
498        return None
499
500    @property
501    def _message_repository(self) -> Optional[MessageRepository]:
502        """
503        The implementation can define a message_repository if it wants debugging logs for HTTP requests
504        """
505        return _NOOP_MESSAGE_REPOSITORY
506
507    def _log_response(self, response: requests.Response) -> None:
508        """
509        Logs the HTTP response using the message repository if it is available.
510
511        Args:
512            response (requests.Response): The HTTP response to log.
513        """
514        if self._message_repository:
515            self._message_repository.log_message(
516                Level.DEBUG,
517                lambda: format_http_message(
518                    response,
519                    "Refresh token",
520                    "Obtains access token",
521                    self._NO_STREAM_NAME,
522                    is_auxiliary=True,
523                    type="AUTH",
524                ),
525            )
526
527    # ----------------
528    # ABSTR METHODS
529    # ----------------
530
531    @abstractmethod
532    def get_token_refresh_endpoint(self) -> Optional[str]:
533        """Returns the endpoint to refresh the access token"""
534
535    @abstractmethod
536    def get_client_id_name(self) -> str:
537        """The client id name to authenticate"""
538
539    @abstractmethod
540    def get_client_id(self) -> str:
541        """The client id to authenticate"""
542
543    @abstractmethod
544    def get_client_secret_name(self) -> str:
545        """The client secret name to authenticate"""
546
547    @abstractmethod
548    def get_client_secret(self) -> str:
549        """The client secret to authenticate"""
550
551    @abstractmethod
552    def get_refresh_token_name(self) -> str:
553        """The refresh token name to authenticate"""
554
555    @abstractmethod
556    def get_refresh_token(self) -> Optional[str]:
557        """The token used to refresh the access token when it expires"""
558
559    @abstractmethod
560    def get_scopes(self) -> List[str]:
561        """List of requested scopes"""
562
563    @abstractmethod
564    def get_token_expiry_date(self) -> AirbyteDateTime:
565        """Expiration date of the access token"""
566
567    @abstractmethod
568    def set_token_expiry_date(self, value: AirbyteDateTime) -> None:
569        """Setter for access token expiration date"""
570
571    @abstractmethod
572    def get_access_token_name(self) -> str:
573        """Field to extract access token from in the response"""
574
575    @abstractmethod
576    def get_expires_in_name(self) -> str:
577        """Returns the expires_in field name"""
578
579    @abstractmethod
580    def get_refresh_request_body(self) -> Mapping[str, Any]:
581        """Returns the request body to set on the refresh request"""
582
583    @abstractmethod
584    def get_refresh_request_headers(self) -> Mapping[str, Any]:
585        """Returns the request headers to set on the refresh request"""
586
587    def should_send_refresh_request_as_query_params(self) -> bool:
588        """Returns `True` if the standard refresh args should be sent on the URL
589        query string instead of in the request body.
590
591        Defaults to `False` so existing authenticators retain their previous
592        behavior (params in body, no query params on the refresh URL). Subclasses
593        can override this to opt into the URL-query-string shape required by OAuth
594        providers like Gong.
595        """
596        return False
597
598    @abstractmethod
599    def get_grant_type(self) -> str:
600        """Returns grant_type specified for requesting access_token"""
601
602    @abstractmethod
603    def get_grant_type_name(self) -> str:
604        """Returns grant_type specified name for requesting access_token"""
605
606    @property
607    @abstractmethod
608    def access_token(self) -> str:
609        """Returns the access token"""
610
611    @access_token.setter
612    @abstractmethod
613    def access_token(self, value: str) -> str:
614        """Setter for the access token"""