airbyte_cdk.sources.declarative.auth.token_provider

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4
  5
  6import datetime
  7import threading
  8from abc import abstractmethod
  9from dataclasses import InitVar, dataclass, field
 10from typing import Any, List, Mapping, Optional, Union
 11
 12import dpath
 13from isodate import Duration
 14
 15from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
 16from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder
 17from airbyte_cdk.sources.declarative.exceptions import ReadException
 18from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
 19from airbyte_cdk.sources.declarative.requesters.requester import Requester
 20from airbyte_cdk.sources.http_logger import format_http_message
 21from airbyte_cdk.sources.message import MessageRepository, NoopMessageRepository
 22from airbyte_cdk.sources.types import Config
 23from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now
 24
 25
 26class TokenProvider:
 27    @abstractmethod
 28    def get_token(self) -> str:
 29        pass
 30
 31
 32@dataclass
 33class SessionTokenProvider(TokenProvider):
 34    login_requester: Requester
 35    session_token_path: List[str]
 36    expiration_duration: Optional[Union[datetime.timedelta, Duration]]
 37    parameters: InitVar[Mapping[str, Any]]
 38    message_repository: MessageRepository = NoopMessageRepository()
 39    decoder: Decoder = field(default_factory=lambda: JsonDecoder(parameters={}))
 40
 41    _next_expiration_time: Optional[AirbyteDateTime] = None
 42    _token: Optional[str] = None
 43
 44    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 45        self._lock = threading.Lock()
 46
 47    def get_token(self) -> str:
 48        self._refresh_if_necessary()
 49        if self._token is None:
 50            raise ReadException("Failed to get session token, token is None")
 51        return self._token
 52
 53    def _refresh_if_necessary(self) -> None:
 54        if self.expiration_duration is None:
 55            # Without an expiry nothing is cached and every call logs in, so there is nothing to lock
 56            self._refresh()
 57        elif self._token_has_expired():
 58            with self._lock:
 59                if self._token_has_expired():
 60                    self._refresh()
 61
 62    def _token_has_expired(self) -> bool:
 63        return self._next_expiration_time is None or self._next_expiration_time < ab_datetime_now()
 64
 65    def _refresh(self) -> None:
 66        response = self.login_requester.send_request(
 67            log_formatter=lambda response: format_http_message(
 68                response,
 69                "Login request",
 70                "Obtains session token",
 71                None,
 72                is_auxiliary=True,
 73                type="AUTH",
 74            ),
 75        )
 76        if response is None:
 77            raise ReadException("Failed to get session token, response got ignored by requester")
 78        session_token = dpath.get(next(self.decoder.decode(response)), self.session_token_path)
 79        # Set the token first: `get_token` reads it without the lock as soon as the expiry is fresh
 80        self._token = session_token  # type: ignore # Returned decoded response will be Mapping and therefore session_token will be str or None
 81        if self.expiration_duration is not None:
 82            self._next_expiration_time = ab_datetime_now() + self.expiration_duration
 83
 84
 85@dataclass
 86class InterpolatedStringTokenProvider(TokenProvider):
 87    """Provides a token by interpolating a string with config values."""
 88
 89    config: Config
 90    api_token: Union[InterpolatedString, str]
 91    parameters: Mapping[str, Any]
 92
 93    def __post_init__(self) -> None:
 94        self._token = InterpolatedString.create(self.api_token, parameters=self.parameters)
 95
 96    def get_token(self) -> str:
 97        return str(self._token.eval(self.config))
 98
 99
100@dataclass
101class InterpolatedSessionTokenProvider(TokenProvider):
102    """Provides a token by interpolating a template with the session token.
103
104    This allows flexible token formatting, such as "Token {{ session_token }}"
105    for Django REST Framework APIs that expect "Authorization: Token <value>".
106    """
107
108    config: Config
109    api_token: Union[InterpolatedString, str]
110    session_token_provider: TokenProvider
111    parameters: Mapping[str, Any]
112
113    def __post_init__(self) -> None:
114        self._token_template = InterpolatedString.create(self.api_token, parameters=self.parameters)
115
116    def get_token(self) -> str:
117        session_token = self.session_token_provider.get_token()
118        return str(self._token_template.eval(self.config, session_token=session_token))