airbyte-cdk 0.0.0.dev0__py3-none-any.whl

airbyte_cdk/sources/declarative/async_job/timer.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+
+
+class Timer:
+    def __init__(self, timeout: timedelta) -> None:
+        self._start_datetime: Optional[datetime] = None
+        self._end_datetime: Optional[datetime] = None
+        self._timeout = timeout
+
+    def start(self) -> None:
+        self._start_datetime = self._now()
+        self._end_datetime = None
+
+    def stop(self) -> None:
+        if self._end_datetime is None:
+            self._end_datetime = self._now()
+
+    def is_started(self) -> bool:
+        return self._start_datetime is not None
+
+    @property
+    def elapsed_time(self) -> Optional[timedelta]:
+        if not self._start_datetime:
+            return None
+
+        end_time = self._end_datetime or self._now()
+        elapsed_period = end_time - self._start_datetime
+        return elapsed_period
+
+    def has_timed_out(self) -> bool:
+        if not self.is_started():
+            return False
+        return self.elapsed_time > self._timeout  # type: ignore  # given the job timer is started, we assume there is an elapsed_period
+
+    @staticmethod
+    def _now() -> datetime:
+        return datetime.now(tz=timezone.utc)

airbyte_cdk/sources/declarative/auth/__init__.py

@@ -0,0 +1,8 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from airbyte_cdk.sources.declarative.auth.jwt import JwtAuthenticator
+from airbyte_cdk.sources.declarative.auth.oauth import DeclarativeOauth2Authenticator
+
+__all__ = ["DeclarativeOauth2Authenticator", "JwtAuthenticator"]

airbyte_cdk/sources/declarative/auth/declarative_authenticator.py

@@ -0,0 +1,42 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from dataclasses import InitVar, dataclass
+from typing import Any, Mapping, Union
+
+from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_token import (
+    AbstractHeaderAuthenticator,
+)
+
+
+@dataclass
+class DeclarativeAuthenticator(AbstractHeaderAuthenticator):
+    """
+    Interface used to associate which authenticators can be used as part of the declarative framework
+    """
+
+    def get_request_params(self) -> Mapping[str, Any]:
+        """HTTP request parameter to add to the requests"""
+        return {}
+
+    def get_request_body_data(self) -> Union[Mapping[str, Any], str]:
+        """Form-encoded body data to set on the requests"""
+        return {}
+
+    def get_request_body_json(self) -> Mapping[str, Any]:
+        """JSON-encoded body data to set on the requests"""
+        return {}
+
+
+@dataclass
+class NoAuth(DeclarativeAuthenticator):
+    parameters: InitVar[Mapping[str, Any]]
+
+    @property
+    def auth_header(self) -> str:
+        return ""
+
+    @property
+    def token(self) -> str:
+        return ""

airbyte_cdk/sources/declarative/auth/jwt.py

@@ -0,0 +1,197 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import base64
+import json
+from dataclasses import InitVar, dataclass
+from datetime import datetime
+from typing import Any, Mapping, Optional, Union
+
+import jwt
+
+from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
+from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean
+from airbyte_cdk.sources.declarative.interpolation.interpolated_mapping import InterpolatedMapping
+from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
+
+
+class JwtAlgorithm(str):
+    """
+    Enum for supported JWT algorithms
+    """
+
+    HS256 = "HS256"
+    HS384 = "HS384"
+    HS512 = "HS512"
+    ES256 = "ES256"
+    ES256K = "ES256K"
+    ES384 = "ES384"
+    ES512 = "ES512"
+    RS256 = "RS256"
+    RS384 = "RS384"
+    RS512 = "RS512"
+    PS256 = "PS256"
+    PS384 = "PS384"
+    PS512 = "PS512"
+    EdDSA = "EdDSA"
+
+
+@dataclass
+class JwtAuthenticator(DeclarativeAuthenticator):
+    """
+    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.
+
+    Attributes:
+        config (Mapping[str, Any]): The user-provided configuration as specified by the source's spec
+        secret_key (Union[InterpolatedString, str]): The secret key used to sign the JWT
+        algorithm (Union[str, JwtAlgorithm]): The algorithm used to sign the JWT
+        token_duration (Optional[int]): The duration in seconds for which the token is valid
+        base64_encode_secret_key (Optional[Union[InterpolatedBoolean, str, bool]]): Whether to base64 encode the secret key
+        header_prefix (Optional[Union[InterpolatedString, str]]): The prefix to add to the Authorization header
+        kid (Optional[Union[InterpolatedString, str]]): The key identifier to be included in the JWT header
+        typ (Optional[Union[InterpolatedString, str]]): The type of the JWT.
+        cty (Optional[Union[InterpolatedString, str]]): The content type of the JWT.
+        iss (Optional[Union[InterpolatedString, str]]): The issuer of the JWT.
+        sub (Optional[Union[InterpolatedString, str]]): The subject of the JWT.
+        aud (Optional[Union[InterpolatedString, str]]): The audience of the JWT.
+        additional_jwt_headers (Optional[Mapping[str, Any]]): Additional headers to include in the JWT.
+        additional_jwt_payload (Optional[Mapping[str, Any]]): Additional payload to include in the JWT.
+    """
+
+    config: Mapping[str, Any]
+    parameters: InitVar[Mapping[str, Any]]
+    secret_key: Union[InterpolatedString, str]
+    algorithm: Union[str, JwtAlgorithm]
+    token_duration: Optional[int]
+    base64_encode_secret_key: Optional[Union[InterpolatedBoolean, str, bool]] = False
+    header_prefix: Optional[Union[InterpolatedString, str]] = None
+    kid: Optional[Union[InterpolatedString, str]] = None
+    typ: Optional[Union[InterpolatedString, str]] = None
+    cty: Optional[Union[InterpolatedString, str]] = None
+    iss: Optional[Union[InterpolatedString, str]] = None
+    sub: Optional[Union[InterpolatedString, str]] = None
+    aud: Optional[Union[InterpolatedString, str]] = None
+    additional_jwt_headers: Optional[Mapping[str, Any]] = None
+    additional_jwt_payload: Optional[Mapping[str, Any]] = None
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        self._secret_key = InterpolatedString.create(self.secret_key, parameters=parameters)
+        self._algorithm = (
+            JwtAlgorithm(self.algorithm) if isinstance(self.algorithm, str) else self.algorithm
+        )
+        self._base64_encode_secret_key = (
+            InterpolatedBoolean(self.base64_encode_secret_key, parameters=parameters)
+            if isinstance(self.base64_encode_secret_key, str)
+            else self.base64_encode_secret_key
+        )
+        self._token_duration = self.token_duration
+        self._header_prefix = (
+            InterpolatedString.create(self.header_prefix, parameters=parameters)
+            if self.header_prefix
+            else None
+        )
+        self._kid = InterpolatedString.create(self.kid, parameters=parameters) if self.kid else None
+        self._typ = InterpolatedString.create(self.typ, parameters=parameters) if self.typ else None
+        self._cty = InterpolatedString.create(self.cty, parameters=parameters) if self.cty else None
+        self._iss = InterpolatedString.create(self.iss, parameters=parameters) if self.iss else None
+        self._sub = InterpolatedString.create(self.sub, parameters=parameters) if self.sub else None
+        self._aud = InterpolatedString.create(self.aud, parameters=parameters) if self.aud else None
+        self._additional_jwt_headers = InterpolatedMapping(
+            self.additional_jwt_headers or {}, parameters=parameters
+        )
+        self._additional_jwt_payload = InterpolatedMapping(
+            self.additional_jwt_payload or {}, parameters=parameters
+        )
+
+    def _get_jwt_headers(self) -> dict[str, Any]:
+        """
+        Builds and returns the headers used when signing the JWT.
+        """
+        headers = self._additional_jwt_headers.eval(self.config, json_loads=json.loads)
+        if any(prop in headers for prop in ["kid", "alg", "typ", "cty"]):
+            raise ValueError(
+                "'kid', 'alg', 'typ', 'cty' are reserved headers and should not be set as part of 'additional_jwt_headers'"
+            )
+
+        if self._kid:
+            headers["kid"] = self._kid.eval(self.config, json_loads=json.loads)
+        if self._typ:
+            headers["typ"] = self._typ.eval(self.config, json_loads=json.loads)
+        if self._cty:
+            headers["cty"] = self._cty.eval(self.config, json_loads=json.loads)
+        headers["alg"] = self._algorithm
+        return headers
+
+    def _get_jwt_payload(self) -> dict[str, Any]:
+        """
+        Builds and returns the payload used when signing the JWT.
+        """
+        now = int(datetime.now().timestamp())
+        exp = now + self._token_duration if isinstance(self._token_duration, int) else now
+        nbf = now
+
+        payload = self._additional_jwt_payload.eval(self.config, json_loads=json.loads)
+        if any(prop in payload for prop in ["iss", "sub", "aud", "iat", "exp", "nbf"]):
+            raise ValueError(
+                "'iss', 'sub', 'aud', 'iat', 'exp', 'nbf' are reserved properties and should not be set as part of 'additional_jwt_payload'"
+            )
+
+        if self._iss:
+            payload["iss"] = self._iss.eval(self.config, json_loads=json.loads)
+        if self._sub:
+            payload["sub"] = self._sub.eval(self.config, json_loads=json.loads)
+        if self._aud:
+            payload["aud"] = self._aud.eval(self.config, json_loads=json.loads)
+
+        payload["iat"] = now
+        payload["exp"] = exp
+        payload["nbf"] = nbf
+        return payload
+
+    def _get_secret_key(self) -> str:
+        """
+        Returns the secret key used to sign the JWT.
+        """
+        secret_key: str = self._secret_key.eval(self.config, json_loads=json.loads)
+        return (
+            base64.b64encode(secret_key.encode()).decode()
+            if self._base64_encode_secret_key
+            else secret_key
+        )
+
+    def _get_signed_token(self) -> Union[str, Any]:
+        """
+        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/
+        """
+        try:
+            return jwt.encode(
+                payload=self._get_jwt_payload(),
+                key=self._get_secret_key(),
+                algorithm=self._algorithm,
+                headers=self._get_jwt_headers(),
+            )
+        except Exception as e:
+            raise ValueError(f"Failed to sign token: {e}")
+
+    def _get_header_prefix(self) -> Union[str, None]:
+        """
+        Returns the header prefix to be used when attaching the token to the request.
+        """
+        return (
+            self._header_prefix.eval(self.config, json_loads=json.loads)
+            if self._header_prefix
+            else None
+        )
+
+    @property
+    def auth_header(self) -> str:
+        return "Authorization"
+
+    @property
+    def token(self) -> str:
+        return (
+            f"{self._get_header_prefix()} {self._get_signed_token()}"
+            if self._get_header_prefix()
+            else self._get_signed_token()
+        )

airbyte_cdk/sources/declarative/auth/oauth.py

@@ -0,0 +1,293 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from dataclasses import InitVar, dataclass, field
+from datetime import datetime, timedelta
+from typing import Any, List, Mapping, MutableMapping, Optional, Union
+
+from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
+from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean
+from airbyte_cdk.sources.declarative.interpolation.interpolated_mapping import InterpolatedMapping
+from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
+from airbyte_cdk.sources.message import MessageRepository, NoopMessageRepository
+from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_oauth import (
+    AbstractOauth2Authenticator,
+)
+from airbyte_cdk.sources.streams.http.requests_native_auth.oauth import (
+    SingleUseRefreshTokenOauth2Authenticator,
+)
+from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now, ab_datetime_parse
+
+
+@dataclass
+class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAuthenticator):
+    """
+    Generates OAuth2.0 access tokens from an OAuth2.0 refresh token and client credentials based on
+    a declarative connector configuration file. Credentials can be defined explicitly or via interpolation
+    at runtime. The generated access token is attached to each request via the Authorization header.
+
+    Attributes:
+        token_refresh_endpoint (Union[InterpolatedString, str]): The endpoint to refresh the access token
+        client_id (Union[InterpolatedString, str]): The client id
+        client_secret (Union[InterpolatedString, str]): Client secret
+        refresh_token (Union[InterpolatedString, str]): The token used to refresh the access token
+        access_token_name (Union[InterpolatedString, str]): THe field to extract access token from in the response
+        expires_in_name (Union[InterpolatedString, str]): The field to extract expires_in from in the response
+        config (Mapping[str, Any]): The user-provided configuration as specified by the source's spec
+        scopes (Optional[List[str]]): The scopes to request
+        token_expiry_date (Optional[Union[InterpolatedString, str]]): The access token expiration date
+        token_expiry_date_format str: format of the datetime; provide it if expires_in is returned in datetime instead of seconds
+        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
+        refresh_request_body (Optional[Mapping[str, Any]]): The request body to send in the refresh request
+        refresh_request_headers (Optional[Mapping[str, Any]]): The request headers to send in the refresh request
+        grant_type: The grant_type to request for access_token. If set to refresh_token, the refresh_token parameter has to be provided
+        message_repository (MessageRepository): the message repository used to emit logs on HTTP requests
+    """
+
+    config: Mapping[str, Any]
+    parameters: InitVar[Mapping[str, Any]]
+    client_id: Optional[Union[InterpolatedString, str]] = None
+    client_secret: Optional[Union[InterpolatedString, str]] = None
+    token_refresh_endpoint: Optional[Union[InterpolatedString, str]] = None
+    refresh_token: Optional[Union[InterpolatedString, str]] = None
+    scopes: Optional[List[str]] = None
+    token_expiry_date: Optional[Union[InterpolatedString, str]] = None
+    _token_expiry_date: Optional[AirbyteDateTime] = field(init=False, repr=False, default=None)
+    token_expiry_date_format: Optional[str] = None
+    token_expiry_is_time_of_expiration: bool = False
+    access_token_name: Union[InterpolatedString, str] = "access_token"
+    access_token_value: Optional[Union[InterpolatedString, str]] = None
+    client_id_name: Union[InterpolatedString, str] = "client_id"
+    client_secret_name: Union[InterpolatedString, str] = "client_secret"
+    expires_in_name: Union[InterpolatedString, str] = "expires_in"
+    refresh_token_name: Union[InterpolatedString, str] = "refresh_token"
+    refresh_request_body: Optional[Mapping[str, Any]] = None
+    refresh_request_headers: Optional[Mapping[str, Any]] = None
+    grant_type_name: Union[InterpolatedString, str] = "grant_type"
+    grant_type: Union[InterpolatedString, str] = "refresh_token"
+    message_repository: MessageRepository = NoopMessageRepository()
+    profile_assertion: Optional[DeclarativeAuthenticator] = None
+    use_profile_assertion: Optional[Union[InterpolatedBoolean, str, bool]] = False
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        super().__init__()
+        if self.token_refresh_endpoint is not None:
+            self._token_refresh_endpoint: Optional[InterpolatedString] = InterpolatedString.create(
+                self.token_refresh_endpoint, parameters=parameters
+            )
+        else:
+            self._token_refresh_endpoint = None
+        self._client_id_name = InterpolatedString.create(self.client_id_name, parameters=parameters)
+        self._client_id = (
+            InterpolatedString.create(self.client_id, parameters=parameters)
+            if self.client_id
+            else self.client_id
+        )
+        self._client_secret_name = InterpolatedString.create(
+            self.client_secret_name, parameters=parameters
+        )
+        self._client_secret = (
+            InterpolatedString.create(self.client_secret, parameters=parameters)
+            if self.client_secret
+            else self.client_secret
+        )
+        self._refresh_token_name = InterpolatedString.create(
+            self.refresh_token_name, parameters=parameters
+        )
+        if self.refresh_token is not None:
+            self._refresh_token: Optional[InterpolatedString] = InterpolatedString.create(
+                self.refresh_token, parameters=parameters
+            )
+        else:
+            self._refresh_token = None
+        self.access_token_name = InterpolatedString.create(
+            self.access_token_name, parameters=parameters
+        )
+        self.expires_in_name = InterpolatedString.create(
+            self.expires_in_name, parameters=parameters
+        )
+        self.grant_type_name = InterpolatedString.create(
+            self.grant_type_name, parameters=parameters
+        )
+        self.grant_type = InterpolatedString.create(
+            "urn:ietf:params:oauth:grant-type:jwt-bearer"
+            if self.use_profile_assertion
+            else self.grant_type,
+            parameters=parameters,
+        )
+        self._refresh_request_body = InterpolatedMapping(
+            self.refresh_request_body or {}, parameters=parameters
+        )
+        self._refresh_request_headers = InterpolatedMapping(
+            self.refresh_request_headers or {}, parameters=parameters
+        )
+        try:
+            if (
+                isinstance(self.token_expiry_date, (int, str))
+                and str(self.token_expiry_date).isdigit()
+            ):
+                self._token_expiry_date = ab_datetime_parse(self.token_expiry_date)
+            else:
+                self._token_expiry_date = (
+                    ab_datetime_parse(
+                        InterpolatedString.create(
+                            self.token_expiry_date, parameters=parameters
+                        ).eval(self.config)
+                    )
+                    if self.token_expiry_date
+                    else ab_datetime_now() - timedelta(days=1)
+                )
+        except ValueError as e:
+            raise ValueError(f"Invalid token expiry date format: {e}")
+        self.use_profile_assertion = (
+            InterpolatedBoolean(self.use_profile_assertion, parameters=parameters)
+            if isinstance(self.use_profile_assertion, str)
+            else self.use_profile_assertion
+        )
+        self.assertion_name = "assertion"
+
+        if self.access_token_value is not None:
+            self._access_token_value = InterpolatedString.create(
+                self.access_token_value, parameters=parameters
+            ).eval(self.config)
+        else:
+            self._access_token_value = None
+
+        self._access_token: Optional[str] = (
+            self._access_token_value if self.access_token_value else None
+        )
+
+        if not self.use_profile_assertion and any(
+            client_creds is None for client_creds in [self.client_id, self.client_secret]
+        ):
+            raise ValueError(
+                "OAuthAuthenticator configuration error: Both 'client_id' and 'client_secret' are required for the "
+                "basic OAuth flow."
+            )
+        if self.profile_assertion is None and self.use_profile_assertion:
+            raise ValueError(
+                "OAuthAuthenticator configuration error: 'profile_assertion' is required when using the profile assertion flow."
+            )
+        if self.get_grant_type() == "refresh_token" and self._refresh_token is None:
+            raise ValueError(
+                "OAuthAuthenticator configuration error: A 'refresh_token' is required when the 'grant_type' is set to 'refresh_token'."
+            )
+
+    def get_token_refresh_endpoint(self) -> Optional[str]:
+        if self._token_refresh_endpoint is not None:
+            refresh_token_endpoint: str = self._token_refresh_endpoint.eval(self.config)
+            if not refresh_token_endpoint:
+                raise ValueError(
+                    "OAuthAuthenticator was unable to evaluate token_refresh_endpoint parameter"
+                )
+            return refresh_token_endpoint
+        return None
+
+    def get_client_id_name(self) -> str:
+        return self._client_id_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
+    def get_client_id(self) -> str:
+        client_id = self._client_id.eval(self.config) if self._client_id else self._client_id
+        if not client_id:
+            raise ValueError("OAuthAuthenticator was unable to evaluate client_id parameter")
+        return client_id  # type: ignore # value will be returned as a string, or an error will be raised
+
+    def get_client_secret_name(self) -> str:
+        return self._client_secret_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
+    def get_client_secret(self) -> str:
+        client_secret = (
+            self._client_secret.eval(self.config) if self._client_secret else self._client_secret
+        )
+        if not client_secret:
+            raise ValueError("OAuthAuthenticator was unable to evaluate client_secret parameter")
+        return client_secret  # type: ignore # value will be returned as a string, or an error will be raised
+
+    def get_refresh_token_name(self) -> str:
+        return self._refresh_token_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
+    def get_refresh_token(self) -> Optional[str]:
+        return None if self._refresh_token is None else str(self._refresh_token.eval(self.config))
+
+    def get_scopes(self) -> List[str]:
+        return self.scopes or []
+
+    def get_access_token_name(self) -> str:
+        return self.access_token_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
+    def get_expires_in_name(self) -> str:
+        return self.expires_in_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
+    def get_grant_type_name(self) -> str:
+        return self.grant_type_name.eval(self.config)  # type: ignore # eval returns a string in this context
+
+    def get_grant_type(self) -> str:
+        return self.grant_type.eval(self.config)  # type: ignore # eval returns a string in this context
+
+    def get_refresh_request_body(self) -> Mapping[str, Any]:
+        return self._refresh_request_body.eval(self.config)
+
+    def get_refresh_request_headers(self) -> Mapping[str, Any]:
+        return self._refresh_request_headers.eval(self.config)
+
+    def get_token_expiry_date(self) -> AirbyteDateTime:
+        if not self._has_access_token_been_initialized():
+            return AirbyteDateTime.from_datetime(datetime.min)
+        return self._token_expiry_date  # type: ignore # _token_expiry_date is an AirbyteDateTime. It is never None despite what mypy thinks
+
+    def _has_access_token_been_initialized(self) -> bool:
+        return self._access_token is not None
+
+    def set_token_expiry_date(self, value: Union[str, int]) -> None:
+        self._token_expiry_date = self._parse_token_expiration_date(value)
+
+    def get_assertion_name(self) -> str:
+        return self.assertion_name
+
+    def get_assertion(self) -> str:
+        if self.profile_assertion is None:
+            raise ValueError("profile_assertion is not set")
+        return self.profile_assertion.token
+
+    def build_refresh_request_body(self) -> Mapping[str, Any]:
+        """
+        Returns the request body to set on the refresh request
+
+        Override to define additional parameters
+        """
+        if self.use_profile_assertion:
+            return {
+                self.get_grant_type_name(): self.get_grant_type(),
+                self.get_assertion_name(): self.get_assertion(),
+            }
+        return super().build_refresh_request_body()
+
+    @property
+    def access_token(self) -> str:
+        if self._access_token is None:
+            raise ValueError("access_token is not set")
+        return self._access_token
+
+    @access_token.setter
+    def access_token(self, value: str) -> None:
+        self._access_token = value
+
+    @property
+    def _message_repository(self) -> MessageRepository:
+        """
+        Overriding AbstractOauth2Authenticator._message_repository to allow for HTTP request logs
+        """
+        return self.message_repository
+
+
+@dataclass
+class DeclarativeSingleUseRefreshTokenOauth2Authenticator(
+    SingleUseRefreshTokenOauth2Authenticator, DeclarativeAuthenticator
+):
+    """
+    Declarative version of SingleUseRefreshTokenOauth2Authenticator which can be used in declarative connectors.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)

airbyte_cdk/sources/declarative/auth/selective_authenticator.py

@@ -0,0 +1,45 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from dataclasses import dataclass
+from typing import Any, List, Mapping
+
+import dpath
+
+from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
+
+
+@dataclass
+class SelectiveAuthenticator(DeclarativeAuthenticator):
+    """Authenticator that selects concrete implementation based on specific config value."""
+
+    config: Mapping[str, Any]
+    authenticators: Mapping[str, DeclarativeAuthenticator]
+    authenticator_selection_path: List[str]
+
+    # returns "DeclarativeAuthenticator", but must return a subtype of "SelectiveAuthenticator"
+    def __new__(  # type: ignore[misc]
+        cls,
+        config: Mapping[str, Any],
+        authenticators: Mapping[str, DeclarativeAuthenticator],
+        authenticator_selection_path: List[str],
+        *arg: Any,
+        **kwargs: Any,
+    ) -> DeclarativeAuthenticator:
+        try:
+            selected_key = str(
+                dpath.get(
+                    config,  # type: ignore[arg-type]  # Dpath wants mutable mapping but doesn't need it.
+                    authenticator_selection_path,
+                )
+            )
+        except KeyError as err:
+            raise ValueError(
+                "The path from `authenticator_selection_path` is not found in the config."
+            ) from err
+
+        try:
+            return authenticators[selected_key]
+        except KeyError as err:
+            raise ValueError(f"The authenticator `{selected_key}` is not found.") from err