airbyte-cdk 0.72.0__py3-none-any.whl → 6.13.1.dev4106__py3-none-any.whl

airbyte_cdk/sources/declarative/auth/declarative_authenticator.py

@@ -5,7 +5,9 @@
 from dataclasses import InitVar, dataclass
 from typing import Any, Mapping, Union
 
-from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_token import AbstractHeaderAuthenticator
+from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_token import (
+    AbstractHeaderAuthenticator,
+)
 
 
 @dataclass

airbyte_cdk/sources/declarative/auth/jwt.py

@@ -0,0 +1,191 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import base64
+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)
+        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)
+        if self._typ:
+            headers["typ"] = self._typ.eval(self.config)
+        if self._cty:
+            headers["cty"] = self._cty.eval(self.config)
+        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)
+        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)
+        if self._sub:
+            payload["sub"] = self._sub.eval(self.config)
+        if self._aud:
+            payload["aud"] = self._aud.eval(self.config)
+        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)
+        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) 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

@@ -6,12 +6,17 @@ from dataclasses import InitVar, dataclass, field
 from typing import Any, List, Mapping, Optional, Union
 
 import pendulum
+
 from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
 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.sources.streams.http.requests_native_auth.abstract_oauth import (
+    AbstractOauth2Authenticator,
+)
+from airbyte_cdk.sources.streams.http.requests_native_auth.oauth import (
+    SingleUseRefreshTokenOauth2Authenticator,
+)
 
 
 @dataclass
@@ -38,11 +43,11 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
         message_repository (MessageRepository): the message repository used to emit logs on HTTP requests
     """
 
-    token_refresh_endpoint: Union[InterpolatedString, str]
     client_id: Union[InterpolatedString, str]
     client_secret: Union[InterpolatedString, str]
     config: Mapping[str, Any]
     parameters: InitVar[Mapping[str, Any]]
+    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
@@ -50,6 +55,7 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
     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
     expires_in_name: Union[InterpolatedString, str] = "expires_in"
     refresh_request_body: Optional[Mapping[str, Any]] = None
     grant_type: Union[InterpolatedString, str] = "refresh_token"
@@ -57,32 +63,64 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         super().__init__()
-        self._token_refresh_endpoint = InterpolatedString.create(self.token_refresh_endpoint, parameters=parameters)
+        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 = InterpolatedString.create(self.client_id, parameters=parameters)
         self._client_secret = InterpolatedString.create(self.client_secret, parameters=parameters)
         if self.refresh_token is not None:
-            self._refresh_token = InterpolatedString.create(self.refresh_token, parameters=parameters)
+            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.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 = InterpolatedString.create(self.grant_type, parameters=parameters)
-        self._refresh_request_body = InterpolatedMapping(self.refresh_request_body or {}, parameters=parameters)
+        self._refresh_request_body = InterpolatedMapping(
+            self.refresh_request_body or {}, parameters=parameters
+        )
         self._token_expiry_date: pendulum.DateTime = (
-            pendulum.parse(InterpolatedString.create(self.token_expiry_date, parameters=parameters).eval(self.config))  # type: ignore # pendulum.parse returns a datetime in this context
+            pendulum.parse(
+                InterpolatedString.create(self.token_expiry_date, parameters=parameters).eval(
+                    self.config
+                )
+            )  # type: ignore # pendulum.parse returns a datetime in this context
             if self.token_expiry_date
             else pendulum.now().subtract(days=1)  # type: ignore # substract does not have type hints
         )
-        self._access_token: Optional[str] = None  # access_token is initialized by a setter
+        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
 
-        if self.get_grant_type() == "refresh_token" and self._refresh_token is None:
-            raise ValueError("OAuthAuthenticator needs a refresh_token parameter if grant_type is set to `refresh_token`")
+        self._access_token: Optional[str] = (
+            self._access_token_value if self.access_token_value else None
+        )
 
-    def get_token_refresh_endpoint(self) -> str:
-        refresh_token: str = self._token_refresh_endpoint.eval(self.config)
-        if not refresh_token:
-            raise ValueError("OAuthAuthenticator was unable to evaluate token_refresh_endpoint parameter")
-        return refresh_token
+        if self.get_grant_type() == "refresh_token" and self._refresh_token is None:
+            raise ValueError(
+                "OAuthAuthenticator needs a refresh_token parameter if 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(self) -> str:
         client_id: str = self._client_id.eval(self.config)
@@ -97,7 +135,7 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
         return client_secret
 
     def get_refresh_token(self) -> Optional[str]:
-        return None if self._refresh_token is None else self._refresh_token.eval(self.config)
+        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 []
@@ -112,7 +150,7 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
         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)  # type: ignore # eval should return a Mapping in this context
+        return self._refresh_request_body.eval(self.config)
 
     def get_token_expiry_date(self) -> pendulum.DateTime:
         return self._token_expiry_date  # type: ignore # _token_expiry_date is a pendulum.DateTime. It is never None despite what mypy thinks
@@ -139,7 +177,9 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
 
 
 @dataclass
-class DeclarativeSingleUseRefreshTokenOauth2Authenticator(SingleUseRefreshTokenOauth2Authenticator, DeclarativeAuthenticator):
+class DeclarativeSingleUseRefreshTokenOauth2Authenticator(
+    SingleUseRefreshTokenOauth2Authenticator, DeclarativeAuthenticator
+):
     """
     Declarative version of SingleUseRefreshTokenOauth2Authenticator which can be used in declarative connectors.
     """

airbyte_cdk/sources/declarative/auth/selective_authenticator.py

@@ -6,6 +6,7 @@ from dataclasses import dataclass
 from typing import Any, List, Mapping
 
 import dpath
+
 from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
 
 
@@ -27,9 +28,16 @@ class SelectiveAuthenticator(DeclarativeAuthenticator):
         **kwargs: Any,
     ) -> DeclarativeAuthenticator:
         try:
-            selected_key = str(dpath.util.get(config, authenticator_selection_path))
+            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
+            raise ValueError(
+                "The path from `authenticator_selection_path` is not found in the config."
+            ) from err
 
         try:
             return authenticators[selected_key]

airbyte_cdk/sources/declarative/auth/token.py

@@ -8,12 +8,16 @@ from dataclasses import InitVar, dataclass
 from typing import Any, Mapping, Union
 
 import requests
+from cachetools import TTLCache, cached
+
 from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
 from airbyte_cdk.sources.declarative.auth.token_provider import TokenProvider
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
-from airbyte_cdk.sources.declarative.requesters.request_option import RequestOption, RequestOptionType
-from airbyte_cdk.sources.declarative.types import Config
-from cachetools import TTLCache, cached
+from airbyte_cdk.sources.declarative.requesters.request_option import (
+    RequestOption,
+    RequestOptionType,
+)
+from airbyte_cdk.sources.types import Config
 
 
 @dataclass
@@ -42,7 +46,9 @@ class ApiKeyAuthenticator(DeclarativeAuthenticator):
     parameters: InitVar[Mapping[str, Any]]
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
-        self._field_name = InterpolatedString.create(self.request_option.field_name, parameters=parameters)
+        self._field_name = InterpolatedString.create(
+            self.request_option.field_name, parameters=parameters
+        )
 
     @property
     def auth_header(self) -> str:
@@ -127,7 +133,9 @@ class BasicHttpAuthenticator(DeclarativeAuthenticator):
 
     @property
     def token(self) -> str:
-        auth_string = f"{self._username.eval(self.config)}:{self._password.eval(self.config)}".encode("utf8")
+        auth_string = (
+            f"{self._username.eval(self.config)}:{self._password.eval(self.config)}".encode("utf8")
+        )
         b64_encoded = base64.b64encode(auth_string).decode("utf8")
         return f"Basic {b64_encoded}"
 
@@ -164,7 +172,9 @@ def get_new_session_token(api_url: str, username: str, password: str, response_k
     )
     response.raise_for_status()
     if not response.ok:
-        raise ConnectionError(f"Failed to retrieve new session token, response code {response.status_code} because {response.reason}")
+        raise ConnectionError(
+            f"Failed to retrieve new session token, response code {response.status_code} because {response.reason}"
+        )
     return str(response.json()[response_key])
 
 
@@ -208,9 +218,13 @@ class LegacySessionTokenAuthenticator(DeclarativeAuthenticator):
         self._api_url = InterpolatedString.create(self.api_url, parameters=parameters)
         self._header = InterpolatedString.create(self.header, parameters=parameters)
         self._session_token = InterpolatedString.create(self.session_token, parameters=parameters)
-        self._session_token_response_key = InterpolatedString.create(self.session_token_response_key, parameters=parameters)
+        self._session_token_response_key = InterpolatedString.create(
+            self.session_token_response_key, parameters=parameters
+        )
         self._login_url = InterpolatedString.create(self.login_url, parameters=parameters)
-        self._validate_session_url = InterpolatedString.create(self.validate_session_url, parameters=parameters)
+        self._validate_session_url = InterpolatedString.create(
+            self.validate_session_url, parameters=parameters
+        )
 
         self.logger = logging.getLogger("airbyte")
 
@@ -232,7 +246,9 @@ class LegacySessionTokenAuthenticator(DeclarativeAuthenticator):
             self.logger.info("Using generated session token by username and password")
             return get_new_session_token(api_url, username, password, session_token_response_key)
 
-        raise ConnectionError("Invalid credentials: session token is not valid or provide username and password")
+        raise ConnectionError(
+            "Invalid credentials: session token is not valid or provide username and password"
+        )
 
     def is_valid_session_token(self) -> bool:
         try:
@@ -251,4 +267,6 @@ class LegacySessionTokenAuthenticator(DeclarativeAuthenticator):
             self.logger.info("Connection check for source is successful.")
             return True
         else:
-            raise ConnectionError(f"Failed to retrieve new session token, response code {response.status_code} because {response.reason}")
+            raise ConnectionError(
+                f"Failed to retrieve new session token, response code {response.status_code} because {response.reason}"
+            )

airbyte_cdk/sources/declarative/auth/token_provider.py

@@ -5,21 +5,22 @@
 
 import datetime
 from abc import abstractmethod
-from dataclasses import InitVar, dataclass
+from dataclasses import InitVar, dataclass, field
 from typing import Any, List, Mapping, Optional, Union
 
-import dpath.util
+import dpath
 import pendulum
+from isodate import Duration
+from pendulum import DateTime
+
 from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
 from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder
 from airbyte_cdk.sources.declarative.exceptions import ReadException
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
 from airbyte_cdk.sources.declarative.requesters.requester import Requester
-from airbyte_cdk.sources.declarative.types import Config
 from airbyte_cdk.sources.http_logger import format_http_message
 from airbyte_cdk.sources.message import MessageRepository, NoopMessageRepository
-from isodate import Duration
-from pendulum import DateTime
+from airbyte_cdk.sources.types import Config
 
 
 class TokenProvider:
@@ -35,8 +36,8 @@ class SessionTokenProvider(TokenProvider):
     expiration_duration: Optional[Union[datetime.timedelta, Duration]]
     parameters: InitVar[Mapping[str, Any]]
     message_repository: MessageRepository = NoopMessageRepository()
+    decoder: Decoder = field(default_factory=lambda: JsonDecoder(parameters={}))
 
-    _decoder: Decoder = JsonDecoder(parameters={})
     _next_expiration_time: Optional[DateTime] = None
     _token: Optional[str] = None
 
@@ -62,10 +63,10 @@ class SessionTokenProvider(TokenProvider):
         )
         if response is None:
             raise ReadException("Failed to get session token, response got ignored by requester")
-        session_token = dpath.util.get(self._decoder.decode(response), self.session_token_path)
+        session_token = dpath.get(next(self.decoder.decode(response)), self.session_token_path)
         if self.expiration_duration is not None:
             self._next_expiration_time = pendulum.now() + self.expiration_duration
-        self._token = session_token
+        self._token = session_token  # type: ignore # Returned decoded response will be Mapping and therefore session_token will be str or None
 
 
 @dataclass

airbyte_cdk/sources/declarative/checks/check_stream.py

@@ -7,8 +7,8 @@ import traceback
 from dataclasses import InitVar, dataclass
 from typing import Any, List, Mapping, Tuple
 
+from airbyte_cdk import AbstractSource
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
-from airbyte_cdk.sources.source import Source
 from airbyte_cdk.sources.streams.http.availability_strategy import HttpAvailabilityStrategy
 
 
@@ -24,25 +24,33 @@ class CheckStream(ConnectionChecker):
     stream_names: List[str]
     parameters: InitVar[Mapping[str, Any]]
 
-    def __post_init__(self, parameters: Mapping[str, Any]):
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._parameters = parameters
 
-    def check_connection(self, source: Source, logger: logging.Logger, config: Mapping[str, Any]) -> Tuple[bool, any]:
-        streams = source.streams(config)
+    def check_connection(
+        self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
+    ) -> Tuple[bool, Any]:
+        streams = source.streams(config=config)
         stream_name_to_stream = {s.name: s for s in streams}
         if len(streams) == 0:
             return False, f"No streams to connect to from source {source}"
         for stream_name in self.stream_names:
             if stream_name not in stream_name_to_stream.keys():
-                raise ValueError(f"{stream_name} is not part of the catalog. Expected one of {stream_name_to_stream.keys()}.")
+                raise ValueError(
+                    f"{stream_name} is not part of the catalog. Expected one of {stream_name_to_stream.keys()}."
+                )
 
             stream = stream_name_to_stream[stream_name]
-            availability_strategy = stream.availability_strategy or HttpAvailabilityStrategy()
+            availability_strategy = HttpAvailabilityStrategy()
             try:
-                stream_is_available, reason = availability_strategy.check_availability(stream, logger, source)
+                stream_is_available, reason = availability_strategy.check_availability(
+                    stream, logger
+                )
                 if not stream_is_available:
                     return False, reason
             except Exception as error:
-                logger.error(f"Encountered an error trying to connect to stream {stream_name}. Error: \n {traceback.format_exc()}")
+                logger.error(
+                    f"Encountered an error trying to connect to stream {stream_name}. Error: \n {traceback.format_exc()}"
+                )
                 return False, f"Unable to connect to stream {stream_name} - {error}"
         return True, None

airbyte_cdk/sources/declarative/checks/connection_checker.py

@@ -6,7 +6,7 @@ import logging
 from abc import ABC, abstractmethod
 from typing import Any, Mapping, Tuple
 
-from airbyte_cdk.sources.source import Source
+from airbyte_cdk import AbstractSource
 
 
 class ConnectionChecker(ABC):
@@ -15,7 +15,9 @@ class ConnectionChecker(ABC):
     """
 
     @abstractmethod
-    def check_connection(self, source: Source, logger: logging.Logger, config: Mapping[str, Any]) -> Tuple[bool, any]:
+    def check_connection(
+        self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
+    ) -> Tuple[bool, Any]:
         """
         Tests if the input configuration can be used to successfully connect to the integration e.g: if a provided Stripe API token can be used to connect
         to the Stripe API.

airbyte_cdk/sources/declarative/concurrency_level/__init__.py

@@ -0,0 +1,7 @@
+#
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+#
+
+from airbyte_cdk.sources.declarative.concurrency_level.concurrency_level import ConcurrencyLevel
+
+__all__ = ["ConcurrencyLevel"]

airbyte_cdk/sources/declarative/concurrency_level/concurrency_level.py

@@ -0,0 +1,50 @@
+#
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+#
+
+from dataclasses import InitVar, dataclass
+from typing import Any, Mapping, Optional, Union
+
+from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
+from airbyte_cdk.sources.types import Config
+
+
+@dataclass
+class ConcurrencyLevel:
+    """
+    Returns the number of worker threads that should be used when syncing concurrent streams in parallel
+
+    Attributes:
+        default_concurrency (Union[int, str]): The hardcoded integer or interpolation of how many worker threads to use during a sync
+        max_concurrency (Optional[int]): The maximum number of worker threads to use when the default_concurrency is exceeded
+    """
+
+    default_concurrency: Union[int, str]
+    max_concurrency: Optional[int]
+    config: Config
+    parameters: InitVar[Mapping[str, Any]]
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        if isinstance(self.default_concurrency, int):
+            self._default_concurrency: Union[int, InterpolatedString] = self.default_concurrency
+        elif "config" in self.default_concurrency and not self.max_concurrency:
+            raise ValueError(
+                "ConcurrencyLevel requires that max_concurrency be defined if the default_concurrency can be used-specified"
+            )
+        else:
+            self._default_concurrency = InterpolatedString.create(
+                self.default_concurrency, parameters=parameters
+            )
+
+    def get_concurrency_level(self) -> int:
+        if isinstance(self._default_concurrency, InterpolatedString):
+            evaluated_default_concurrency = self._default_concurrency.eval(config=self.config)
+            if not isinstance(evaluated_default_concurrency, int):
+                raise ValueError("default_concurrency did not evaluate to an integer")
+            return (
+                min(evaluated_default_concurrency, self.max_concurrency)
+                if self.max_concurrency
+                else evaluated_default_concurrency
+            )
+        else:
+            return self._default_concurrency