airbyte-cdk 6.31.1__py3-none-any.whl → 6.31.2.dev0__py3-none-any.whl

airbyte_cdk/cli/source_declarative_manifest/_run.py

@@ -21,6 +21,7 @@ import pkgutil
 import sys
 import traceback
 from collections.abc import Mapping
+from datetime import datetime
 from pathlib import Path
 from typing import Any, cast
 
@@ -43,7 +44,6 @@ from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
 )
 from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource
 from airbyte_cdk.sources.source import TState
-from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
 
 
 class SourceLocalYaml(YamlDeclarativeSource):
@@ -101,7 +101,7 @@ def _get_local_yaml_source(args: list[str]) -> SourceLocalYaml:
                         type=Type.TRACE,
                         trace=AirbyteTraceMessage(
                             type=TraceType.ERROR,
-                            emitted_at=ab_datetime_now().to_epoch_millis(),
+                            emitted_at=int(datetime.now().timestamp() * 1000),
                             error=AirbyteErrorTraceMessage(
                                 message=f"Error starting the sync. This could be due to an invalid configuration or catalog. Please contact Support for assistance. Error: {error}",
                                 stack_trace=traceback.format_exc(),
@@ -171,12 +171,6 @@ def create_declarative_source(
                 "Invalid config: `__injected_declarative_manifest` should be provided at the root "
                 f"of the config but config only has keys: {list(config.keys() if config else [])}"
             )
-        if not isinstance(config["__injected_declarative_manifest"], dict):
-            raise ValueError(
-                "Invalid config: `__injected_declarative_manifest` should be a dictionary, "
-                f"but got type: {type(config['__injected_declarative_manifest'])}"
-            )
-
         return ConcurrentDeclarativeSource(
             config=config,
             catalog=catalog,
@@ -191,7 +185,7 @@ def create_declarative_source(
                         type=Type.TRACE,
                         trace=AirbyteTraceMessage(
                             type=TraceType.ERROR,
-                            emitted_at=ab_datetime_now().to_epoch_millis(),
+                            emitted_at=int(datetime.now().timestamp() * 1000),
                             error=AirbyteErrorTraceMessage(
                                 message=f"Error starting the sync. This could be due to an invalid configuration or catalog. Please contact Support for assistance. Error: {error}",
                                 stack_trace=traceback.format_exc(),

airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -3,6 +3,7 @@
 #
 
 import dataclasses
+from datetime import datetime
 from typing import Any, List, Mapping
 
 from airbyte_cdk.connector_builder.message_grouper import MessageGrouper
@@ -20,7 +21,6 @@ from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
 )
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
-from airbyte_cdk.utils.datetime_helpers import ab_datetime_now
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE = 5
@@ -52,7 +52,6 @@ def get_limits(config: Mapping[str, Any]) -> TestReadLimits:
 def create_source(config: Mapping[str, Any], limits: TestReadLimits) -> ManifestDeclarativeSource:
     manifest = config["__injected_declarative_manifest"]
     return ManifestDeclarativeSource(
-        config=config,
         emit_connector_builder_messages=True,
         source_config=manifest,
         component_factory=ModelToComponentFactory(
@@ -114,4 +113,4 @@ def resolve_manifest(source: ManifestDeclarativeSource) -> AirbyteMessage:
 
 
 def _emitted_at() -> int:
-    return ab_datetime_now().to_epoch_millis()
+    return int(datetime.now().timestamp()) * 1000

airbyte_cdk/sources/declarative/async_job/job_orchestrator.py

@@ -482,16 +482,16 @@ class AsyncJobOrchestrator:
             and exception.failure_type == FailureType.config_error
         )
 
-    def fetch_records(self, async_jobs: Iterable[AsyncJob]) -> Iterable[Mapping[str, Any]]:
+    def fetch_records(self, partition: AsyncPartition) -> Iterable[Mapping[str, Any]]:
         """
-        Fetches records from the given jobs.
+        Fetches records from the given partition's jobs.
 
         Args:
-            async_jobs Iterable[AsyncJob]: The list of AsyncJobs.
+            partition (AsyncPartition): The partition containing the jobs.
 
         Yields:
             Iterable[Mapping[str, Any]]: The fetched records from the jobs.
         """
-        for job in async_jobs:
+        for job in partition.jobs:
             yield from self._job_repository.fetch_records(job)
             self._job_repository.delete(job)

airbyte_cdk/sources/declarative/auth/jwt.py

@@ -3,7 +3,6 @@
 #
 
 import base64
-import json
 from dataclasses import InitVar, dataclass
 from datetime import datetime
 from typing import Any, Mapping, Optional, Union
@@ -105,21 +104,21 @@ class JwtAuthenticator(DeclarativeAuthenticator):
         )
 
     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)
+        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, json_loads=json.loads)
+            headers["kid"] = self._kid.eval(self.config)
         if self._typ:
-            headers["typ"] = self._typ.eval(self.config, json_loads=json.loads)
+            headers["typ"] = self._typ.eval(self.config)
         if self._cty:
-            headers["cty"] = self._cty.eval(self.config, json_loads=json.loads)
+            headers["cty"] = self._cty.eval(self.config)
         headers["alg"] = self._algorithm
         return headers
 
@@ -131,19 +130,18 @@ class JwtAuthenticator(DeclarativeAuthenticator):
         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)
+        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, json_loads=json.loads)
+            payload["iss"] = self._iss.eval(self.config)
         if self._sub:
-            payload["sub"] = self._sub.eval(self.config, json_loads=json.loads)
+            payload["sub"] = self._sub.eval(self.config)
         if self._aud:
-            payload["aud"] = self._aud.eval(self.config, json_loads=json.loads)
-
+            payload["aud"] = self._aud.eval(self.config)
         payload["iat"] = now
         payload["exp"] = exp
         payload["nbf"] = nbf
@@ -153,7 +151,7 @@ class JwtAuthenticator(DeclarativeAuthenticator):
         """
         Returns the secret key used to sign the JWT.
         """
-        secret_key: str = self._secret_key.eval(self.config, json_loads=json.loads)
+        secret_key: str = self._secret_key.eval(self.config)
         return (
             base64.b64encode(secret_key.encode()).decode()
             if self._base64_encode_secret_key
@@ -178,11 +176,7 @@ class JwtAuthenticator(DeclarativeAuthenticator):
         """
         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
-        )
+        return self._header_prefix.eval(self.config) if self._header_prefix else None
 
     @property
     def auth_header(self) -> str:

airbyte_cdk/sources/declarative/auth/oauth.py

@@ -3,11 +3,11 @@
 #
 
 from dataclasses import InitVar, dataclass, field
-from datetime import timedelta
-from typing import Any, List, Mapping, MutableMapping, Optional, Union
+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_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
@@ -17,7 +17,6 @@ from airbyte_cdk.sources.streams.http.requests_native_auth.abstract_oauth import
 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
@@ -45,15 +44,15 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
         message_repository (MessageRepository): the message repository used to emit logs on HTTP requests
     """
 
+    client_id: Union[InterpolatedString, str]
+    client_secret: Union[InterpolatedString, str]
     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: Optional[pendulum.DateTime] = 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"
@@ -67,8 +66,6 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
     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__()
@@ -79,19 +76,11 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
         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_id = InterpolatedString.create(self.client_id, parameters=parameters)
         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._client_secret = InterpolatedString.create(self.client_secret, parameters=parameters)
         self._refresh_token_name = InterpolatedString.create(
             self.refresh_token_name, parameters=parameters
         )
@@ -110,43 +99,22 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
         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.grant_type = InterpolatedString.create(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)
+        self._token_expiry_date: pendulum.DateTime = (
+            pendulum.parse(
+                InterpolatedString.create(self.token_expiry_date, parameters=parameters).eval(
+                    self.config
                 )
-        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
+            )  # 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.assertion_name = "assertion"
-
         if self.access_token_value is not None:
             self._access_token_value = InterpolatedString.create(
                 self.access_token_value, parameters=parameters
@@ -158,20 +126,9 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
             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'."
+                "OAuthAuthenticator needs a refresh_token parameter if grant_type is set to `refresh_token`"
             )
 
     def get_token_refresh_endpoint(self) -> Optional[str]:
@@ -188,21 +145,19 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
         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
+        client_id: str = self._client_id.eval(self.config)
         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
+        return client_id
 
     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
-        )
+        client_secret: str = self._client_secret.eval(self.config)
         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
+        return client_secret
 
     def get_refresh_token_name(self) -> str:
         return self._refresh_token_name.eval(self.config)  # type: ignore # eval returns a string in this context
@@ -231,33 +186,12 @@ class DeclarativeOauth2Authenticator(AbstractOauth2Authenticator, DeclarativeAut
     def get_refresh_request_headers(self) -> Mapping[str, Any]:
         return self._refresh_request_headers.eval(self.config)
 
-    def get_token_expiry_date(self) -> AirbyteDateTime:
-        return self._token_expiry_date  # type: ignore # _token_expiry_date is an AirbyteDateTime. It is never None despite what mypy thinks
+    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
 
     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:

airbyte_cdk/sources/declarative/auth/token.py

@@ -5,7 +5,7 @@
 import base64
 import logging
 from dataclasses import InitVar, dataclass
-from typing import Any, Mapping, Union
+from typing import Any, Mapping, MutableMapping, Union
 
 import requests
 from cachetools import TTLCache, cached
@@ -45,11 +45,6 @@ class ApiKeyAuthenticator(DeclarativeAuthenticator):
     config: Config
     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
-        )
-
     @property
     def auth_header(self) -> str:
         options = self._get_request_options(RequestOptionType.header)
@@ -60,9 +55,9 @@ class ApiKeyAuthenticator(DeclarativeAuthenticator):
         return self.token_provider.get_token()
 
     def _get_request_options(self, option_type: RequestOptionType) -> Mapping[str, Any]:
-        options = {}
+        options: MutableMapping[str, Any] = {}
         if self.request_option.inject_into == option_type:
-            options[self._field_name.eval(self.config)] = self.token
+            self.request_option.inject_into_request(options, self.token, self.config)
         return options
 
     def get_request_params(self) -> Mapping[str, Any]:

airbyte_cdk/sources/declarative/auth/token_provider.py

@@ -9,7 +9,9 @@ from dataclasses import InitVar, dataclass, field
 from typing import Any, List, Mapping, Optional, Union
 
 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
@@ -19,7 +21,6 @@ from airbyte_cdk.sources.declarative.requesters.requester import Requester
 from airbyte_cdk.sources.http_logger import format_http_message
 from airbyte_cdk.sources.message import MessageRepository, NoopMessageRepository
 from airbyte_cdk.sources.types import Config
-from airbyte_cdk.utils.datetime_helpers import AirbyteDateTime, ab_datetime_now
 
 
 class TokenProvider:
@@ -37,7 +38,7 @@ class SessionTokenProvider(TokenProvider):
     message_repository: MessageRepository = NoopMessageRepository()
     decoder: Decoder = field(default_factory=lambda: JsonDecoder(parameters={}))
 
-    _next_expiration_time: Optional[AirbyteDateTime] = None
+    _next_expiration_time: Optional[DateTime] = None
     _token: Optional[str] = None
 
     def get_token(self) -> str:
@@ -47,7 +48,7 @@ class SessionTokenProvider(TokenProvider):
         return self._token
 
     def _refresh_if_necessary(self) -> None:
-        if self._next_expiration_time is None or self._next_expiration_time < ab_datetime_now():
+        if self._next_expiration_time is None or self._next_expiration_time < pendulum.now():
             self._refresh()
 
     def _refresh(self) -> None:
@@ -64,7 +65,7 @@ class SessionTokenProvider(TokenProvider):
             raise ReadException("Failed to get session token, response got ignored by requester")
         session_token = dpath.get(next(self.decoder.decode(response)), self.session_token_path)
         if self.expiration_duration is not None:
-            self._next_expiration_time = ab_datetime_now() + self.expiration_duration
+            self._next_expiration_time = pendulum.now() + self.expiration_duration
         self._token = session_token  # type: ignore # Returned decoded response will be Mapping and therefore session_token will be str or None
 
 

airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py

@@ -21,12 +21,8 @@ class CheckDynamicStream(ConnectionChecker):
         stream_count (int): numbers of streams to check
     """
 
-    # TODO: Add field stream_names to check_connection for static streams
-    #  https://github.com/airbytehq/airbyte-python-cdk/pull/293#discussion_r1934933483
-
     stream_count: int
     parameters: InitVar[Mapping[str, Any]]
-    use_check_availability: bool = True
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._parameters = parameters
@@ -35,27 +31,21 @@ class CheckDynamicStream(ConnectionChecker):
         self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
     ) -> Tuple[bool, Any]:
         streams = source.streams(config=config)
-
         if len(streams) == 0:
             return False, f"No streams to connect to from source {source}"
-        if not self.use_check_availability:
-            return True, None
-
-        availability_strategy = HttpAvailabilityStrategy()
 
-        try:
-            for stream in streams[: min(self.stream_count, len(streams))]:
+        for stream_index in range(min(self.stream_count, len(streams))):
+            stream = streams[stream_index]
+            availability_strategy = HttpAvailabilityStrategy()
+            try:
                 stream_is_available, reason = availability_strategy.check_availability(
                     stream, logger
                 )
                 if not stream_is_available:
-                    logger.warning(f"Stream {stream.name} is not available: {reason}")
                     return False, reason
-        except Exception as error:
-            error_message = (
-                f"Encountered an error trying to connect to stream {stream.name}. Error: {error}"
-            )
-            logger.error(error_message, exc_info=True)
-            return False, error_message
-
+            except Exception as error:
+                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