airbyte-cdk 6.30.0.dev0__tar.gz → 6.31.1__tar.gz

{airbyte_cdk-6.30.0.dev0 → airbyte_cdk-6.31.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 6.30.0.dev0
+Version: 6.31.1
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://airbyte.com
 License: MIT

{airbyte_cdk-6.30.0.dev0 → airbyte_cdk-6.31.1}/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, partition: AsyncPartition) -> Iterable[Mapping[str, Any]]:
+    def fetch_records(self, async_jobs: Iterable[AsyncJob]) -> Iterable[Mapping[str, Any]]:
         """
-        Fetches records from the given partition's jobs.
+        Fetches records from the given jobs.
 
         Args:
-            partition (AsyncPartition): The partition containing the jobs.
+            async_jobs Iterable[AsyncJob]: The list of AsyncJobs.
 
         Yields:
             Iterable[Mapping[str, Any]]: The fetched records from the jobs.
         """
-        for job in partition.jobs:
+        for job in async_jobs:
             yield from self._job_repository.fetch_records(job)
             self._job_repository.delete(job)

{airbyte_cdk-6.30.0.dev0 → airbyte_cdk-6.31.1}/airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -19,6 +19,7 @@ from airbyte_cdk.sources.declarative.extractors import RecordSelector
 from airbyte_cdk.sources.declarative.extractors.record_filter import (
     ClientSideIncrementalRecordFilterDecorator,
 )
+from airbyte_cdk.sources.declarative.incremental import ConcurrentPerPartitionCursor
 from airbyte_cdk.sources.declarative.incremental.datetime_based_cursor import DatetimeBasedCursor
 from airbyte_cdk.sources.declarative.incremental.per_partition_with_global import (
     PerPartitionWithGlobalCursor,
@@ -231,7 +232,7 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                     ):
                         cursor = declarative_stream.retriever.stream_slicer.stream_slicer
 
-                        if not isinstance(cursor, ConcurrentCursor):
+                        if not isinstance(cursor, ConcurrentCursor | ConcurrentPerPartitionCursor):
                             # This should never happen since we instantiate ConcurrentCursor in
                             # model_to_component_factory.py
                             raise ValueError(

{airbyte_cdk-6.30.0.dev0 → airbyte_cdk-6.31.1}/airbyte_cdk/sources/declarative/declarative_stream.py

@@ -138,7 +138,9 @@ class DeclarativeStream(Stream):
         """
         :param: stream_state We knowingly avoid using stream_state as we want cursors to manage their own state.
         """
-        if stream_slice is None or stream_slice == {}:
+        if stream_slice is None or (
+            not isinstance(stream_slice, StreamSlice) and stream_slice == {}
+        ):
             # As the parameter is Optional, many would just call `read_records(sync_mode)` during testing without specifying the field
             # As part of the declarative model without custom components, this should never happen as the CDK would wire up a
             # SinglePartitionRouter that would create this StreamSlice properly

{airbyte_cdk-6.30.0.dev0 → airbyte_cdk-6.31.1}/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -1656,7 +1656,7 @@ class ModelToComponentFactory:
     ) -> Optional[PartitionRouter]:
         if (
             hasattr(model, "partition_router")
-            and isinstance(model, SimpleRetrieverModel)
+            and isinstance(model, SimpleRetrieverModel | AsyncRetrieverModel)
             and model.partition_router
         ):
             stream_slicer_model = model.partition_router
@@ -1690,6 +1690,31 @@ class ModelToComponentFactory:
         stream_slicer = self._build_stream_slicer_from_partition_router(model.retriever, config)
 
         if model.incremental_sync and stream_slicer:
+            if model.retriever.type == "AsyncRetriever":
+                if model.incremental_sync.type != "DatetimeBasedCursor":
+                    # We are currently in a transition to the Concurrent CDK and AsyncRetriever can only work with the support or unordered slices (for example, when we trigger reports for January and February, the report in February can be completed first). Once we have support for custom concurrent cursor or have a new implementation available in the CDK, we can enable more cursors here.
+                    raise ValueError(
+                        "AsyncRetriever with cursor other than DatetimeBasedCursor is not supported yet"
+                    )
+                if stream_slicer:
+                    return self.create_concurrent_cursor_from_perpartition_cursor(  # type: ignore # This is a known issue that we are creating and returning a ConcurrentCursor which does not technically implement the (low-code) StreamSlicer. However, (low-code) StreamSlicer and ConcurrentCursor both implement StreamSlicer.stream_slices() which is the primary method needed for checkpointing
+                        state_manager=self._connector_state_manager,
+                        model_type=DatetimeBasedCursorModel,
+                        component_definition=model.incremental_sync.__dict__,
+                        stream_name=model.name or "",
+                        stream_namespace=None,
+                        config=config or {},
+                        stream_state={},
+                        partition_router=stream_slicer,
+                    )
+                return self.create_concurrent_cursor_from_datetime_based_cursor(  # type: ignore # This is a known issue that we are creating and returning a ConcurrentCursor which does not technically implement the (low-code) StreamSlicer. However, (low-code) StreamSlicer and ConcurrentCursor both implement StreamSlicer.stream_slices() which is the primary method needed for checkpointing
+                    model_type=DatetimeBasedCursorModel,
+                    component_definition=model.incremental_sync.__dict__,
+                    stream_name=model.name or "",
+                    stream_namespace=None,
+                    config=config or {},
+                )
+
             incremental_sync_model = model.incremental_sync
             if (
                 hasattr(incremental_sync_model, "global_substream_cursor")

{airbyte_cdk-6.30.0.dev0 → airbyte_cdk-6.31.1}/airbyte_cdk/sources/declarative/partition_routers/async_job_partition_router.py

@@ -4,9 +4,9 @@ from dataclasses import InitVar, dataclass, field
 from typing import Any, Callable, Iterable, Mapping, Optional
 
 from airbyte_cdk.models import FailureType
+from airbyte_cdk.sources.declarative.async_job.job import AsyncJob
 from airbyte_cdk.sources.declarative.async_job.job_orchestrator import (
     AsyncJobOrchestrator,
-    AsyncPartition,
 )
 from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import (
     SinglePartitionRouter,
@@ -42,12 +42,12 @@ class AsyncJobPartitionRouter(StreamSlicer):
 
         for completed_partition in self._job_orchestrator.create_and_get_completed_partitions():
             yield StreamSlice(
-                partition=dict(completed_partition.stream_slice.partition)
-                | {"partition": completed_partition},
+                partition=dict(completed_partition.stream_slice.partition),
                 cursor_slice=completed_partition.stream_slice.cursor_slice,
+                extra_fields={"jobs": list(completed_partition.jobs)},
             )
 
-    def fetch_records(self, partition: AsyncPartition) -> Iterable[Mapping[str, Any]]:
+    def fetch_records(self, async_jobs: Iterable[AsyncJob]) -> Iterable[Mapping[str, Any]]:
         """
         This method of fetching records extends beyond what a PartitionRouter/StreamSlicer should
         be responsible for. However, this was added in because the JobOrchestrator is required to
@@ -62,4 +62,4 @@ class AsyncJobPartitionRouter(StreamSlicer):
                 failure_type=FailureType.system_error,
             )
 
-        return self._job_orchestrator.fetch_records(partition=partition)
+        return self._job_orchestrator.fetch_records(async_jobs=async_jobs)

{airbyte_cdk-6.30.0.dev0 → airbyte_cdk-6.31.1}/airbyte_cdk/sources/declarative/retrievers/async_retriever.py

@@ -6,7 +6,7 @@ from typing import Any, Iterable, Mapping, Optional
 
 from typing_extensions import deprecated
 
-from airbyte_cdk.models import FailureType
+from airbyte_cdk.sources.declarative.async_job.job import AsyncJob
 from airbyte_cdk.sources.declarative.async_job.job_orchestrator import AsyncPartition
 from airbyte_cdk.sources.declarative.extractors.record_selector import RecordSelector
 from airbyte_cdk.sources.declarative.partition_routers.async_job_partition_router import (
@@ -16,7 +16,6 @@ from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever
 from airbyte_cdk.sources.source import ExperimentalClassWarning
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
-from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
 @deprecated(
@@ -57,9 +56,9 @@ class AsyncRetriever(Retriever):
 
         return self.state
 
-    def _validate_and_get_stream_slice_partition(
+    def _validate_and_get_stream_slice_jobs(
         self, stream_slice: Optional[StreamSlice] = None
-    ) -> AsyncPartition:
+    ) -> Iterable[AsyncJob]:
         """
         Validates the stream_slice argument and returns the partition from it.
 
@@ -73,12 +72,7 @@ class AsyncRetriever(Retriever):
             AirbyteTracedException: If the stream_slice is not an instance of StreamSlice or if the partition is not present in the stream_slice.
 
         """
-        if not isinstance(stream_slice, StreamSlice) or "partition" not in stream_slice.partition:
-            raise AirbyteTracedException(
-                message="Invalid arguments to AsyncRetriever.read_records: stream_slice is not optional. Please contact Airbyte Support",
-                failure_type=FailureType.system_error,
-            )
-        return stream_slice["partition"]  # type: ignore  # stream_slice["partition"] has been added as an AsyncPartition as part of stream_slices
+        return stream_slice.extra_fields.get("jobs", []) if stream_slice else []
 
     def stream_slices(self) -> Iterable[Optional[StreamSlice]]:
         return self.stream_slicer.stream_slices()
@@ -89,8 +83,8 @@ class AsyncRetriever(Retriever):
         stream_slice: Optional[StreamSlice] = None,
     ) -> Iterable[StreamData]:
         stream_state: StreamState = self._get_stream_state()
-        partition: AsyncPartition = self._validate_and_get_stream_slice_partition(stream_slice)
-        records: Iterable[Mapping[str, Any]] = self.stream_slicer.fetch_records(partition)
+        jobs: Iterable[AsyncJob] = self._validate_and_get_stream_slice_jobs(stream_slice)
+        records: Iterable[Mapping[str, Any]] = self.stream_slicer.fetch_records(jobs)
 
         yield from self.record_selector.filter_and_transform(
             all_data=records,

{airbyte_cdk-6.30.0.dev0 → airbyte_cdk-6.31.1}/airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py

@@ -25,6 +25,13 @@ logger = logging.getLogger("airbyte")
 _NOOP_MESSAGE_REPOSITORY = NoopMessageRepository()
 
 
+class ResponseKeysMaxRecurtionReached(AirbyteTracedException):
+    """
+    Raised when the max level of recursion is reached, when trying to
+    find-and-get the target key, during the `_make_handled_request`
+    """
+
+
 class AbstractOauth2Authenticator(AuthBase):
     """
     Abstract class for an OAuth authenticators that implements the OAuth token refresh flow. The authenticator
@@ -53,15 +60,31 @@ class AbstractOauth2Authenticator(AuthBase):
         request.headers.update(self.get_auth_header())
         return request
 
+    @property
+    def _is_access_token_flow(self) -> bool:
+        return self.get_token_refresh_endpoint() is None and self.access_token is not None
+
+    @property
+    def token_expiry_is_time_of_expiration(self) -> bool:
+        """
+        Indicates that the Token Expiry returns the date until which the token will be valid, not the amount of time it will be valid.
+        """
+
+        return False
+
+    @property
+    def token_expiry_date_format(self) -> Optional[str]:
+        """
+        Format of the datetime; exists it if expires_in is returned as the expiration datetime instead of seconds until it expires
+        """
+
+        return None
+
     def get_auth_header(self) -> Mapping[str, Any]:
         """HTTP header to set on the requests"""
         token = self.access_token if self._is_access_token_flow else self.get_access_token()
         return {"Authorization": f"Bearer {token}"}
 
-    @property
-    def _is_access_token_flow(self) -> bool:
-        return self.get_token_refresh_endpoint() is None and self.access_token is not None
-
     def get_access_token(self) -> str:
         """Returns the access token"""
         if self.token_has_expired():
@@ -107,9 +130,39 @@ class AbstractOauth2Authenticator(AuthBase):
         headers = self.get_refresh_request_headers()
         return headers if headers else None
 
+    def refresh_access_token(self) -> Tuple[str, Union[str, int]]:
+        """
+        Returns the refresh token and its expiration datetime
+
+        :return: a tuple of (access_token, token_lifespan)
+        """
+        response_json = self._make_handled_request()
+        self._ensure_access_token_in_response(response_json)
+
+        return (
+            self._extract_access_token(response_json),
+            self._extract_token_expiry_date(response_json),
+        )
+
+    # ----------------
+    # PRIVATE METHODS
+    # ----------------
+
     def _wrap_refresh_token_exception(
         self, exception: requests.exceptions.RequestException
     ) -> bool:
+        """
+        Wraps and handles exceptions that occur during the refresh token process.
+
+        This method checks if the provided exception is related to a refresh token error
+        by examining the response status code and specific error content.
+
+        Args:
+            exception (requests.exceptions.RequestException): The exception raised during the request.
+
+        Returns:
+            bool: True if the exception is related to a refresh token error, False otherwise.
+        """
         try:
             if exception.response is not None:
                 exception_content = exception.response.json()
@@ -131,7 +184,24 @@ class AbstractOauth2Authenticator(AuthBase):
         ),
         max_time=300,
     )
-    def _get_refresh_access_token_response(self) -> Any:
+    def _make_handled_request(self) -> Any:
+        """
+        Makes a handled HTTP request to refresh an OAuth token.
+
+        This method sends a POST request to the token refresh endpoint with the necessary
+        headers and body to obtain a new access token. It handles various exceptions that
+        may occur during the request and logs the response for troubleshooting purposes.
+
+        Returns:
+            Mapping[str, Any]: The JSON response from the token refresh endpoint.
+
+        Raises:
+            DefaultBackoffException: If the response status code is 429 (Too Many Requests)
+                                     or any 5xx server error.
+            AirbyteTracedException: If the refresh token is invalid or expired, prompting
+                                    re-authentication.
+            Exception: For any other exceptions that occur during the request.
+        """
         try:
             response = requests.request(
                 method="POST",
@@ -139,22 +209,10 @@ class AbstractOauth2Authenticator(AuthBase):
                 data=self.build_refresh_request_body(),
                 headers=self.build_refresh_request_headers(),
             )
-            if response.ok:
-                response_json = response.json()
-                # Add the access token to the list of secrets so it is replaced before logging the response
-                # An argument could be made to remove the prevous access key from the list of secrets, but unmasking values seems like a security incident waiting to happen...
-                access_key = response_json.get(self.get_access_token_name())
-                if not access_key:
-                    raise Exception(
-                        "Token refresh API response was missing access token {self.get_access_token_name()}"
-                    )
-                add_to_secrets(access_key)
-                self._log_response(response)
-                return response_json
-            else:
-                # log the response even if the request failed for troubleshooting purposes
-                self._log_response(response)
-                response.raise_for_status()
+            # log the response even if the request failed for troubleshooting purposes
+            self._log_response(response)
+            response.raise_for_status()
+            return response.json()
         except requests.exceptions.RequestException as e:
             if e.response is not None:
                 if e.response.status_code == 429 or e.response.status_code >= 500:
@@ -168,17 +226,34 @@ class AbstractOauth2Authenticator(AuthBase):
         except Exception as e:
             raise Exception(f"Error while refreshing access token: {e}") from e
 
-    def refresh_access_token(self) -> Tuple[str, Union[str, int]]:
+    def _ensure_access_token_in_response(self, response_data: Mapping[str, Any]) -> None:
         """
-        Returns the refresh token and its expiration datetime
+        Ensures that the access token is present in the response data.
 
-        :return: a tuple of (access_token, token_lifespan)
-        """
-        response_json = self._get_refresh_access_token_response()
+        This method attempts to extract the access token from the provided response data.
+        If the access token is not found, it raises an exception indicating that the token
+        refresh API response was missing the access token. If the access token is found,
+        it adds the token to the list of secrets to ensure it is replaced before logging
+        the response.
+
+        Args:
+            response_data (Mapping[str, Any]): The response data from which to extract the access token.
 
-        return response_json[self.get_access_token_name()], response_json[
-            self.get_expires_in_name()
-        ]
+        Raises:
+            Exception: If the access token is not found in the response data.
+            ResponseKeysMaxRecurtionReached: If the maximum recursion depth is reached while extracting the access token.
+        """
+        try:
+            access_key = self._extract_access_token(response_data)
+            if not access_key:
+                raise Exception(
+                    "Token refresh API response was missing access token {self.get_access_token_name()}"
+                )
+            # Add the access token to the list of secrets so it is replaced before logging the response
+            # An argument could be made to remove the prevous access key from the list of secrets, but unmasking values seems like a security incident waiting to happen...
+            add_to_secrets(access_key)
+        except ResponseKeysMaxRecurtionReached as e:
+            raise e
 
     def _parse_token_expiration_date(self, value: Union[str, int]) -> AirbyteDateTime:
         """
@@ -206,22 +281,125 @@ class AbstractOauth2Authenticator(AuthBase):
                     f"Invalid expires_in value: {value}. Expected number of seconds when no format specified."
                 )
 
-    @property
-    def token_expiry_is_time_of_expiration(self) -> bool:
+    def _extract_access_token(self, response_data: Mapping[str, Any]) -> Any:
         """
-        Indicates that the Token Expiry returns the date until which the token will be valid, not the amount of time it will be valid.
+        Extracts the access token from the given response data.
+
+        Args:
+            response_data (Mapping[str, Any]): The response data from which to extract the access token.
+
+        Returns:
+            str: The extracted access token.
         """
+        return self._find_and_get_value_from_response(response_data, self.get_access_token_name())
 
-        return False
+    def _extract_refresh_token(self, response_data: Mapping[str, Any]) -> Any:
+        """
+        Extracts the refresh token from the given response data.
 
-    @property
-    def token_expiry_date_format(self) -> Optional[str]:
+        Args:
+            response_data (Mapping[str, Any]): The response data from which to extract the refresh token.
+
+        Returns:
+            str: The extracted refresh token.
         """
-        Format of the datetime; exists it if expires_in is returned as the expiration datetime instead of seconds until it expires
+        return self._find_and_get_value_from_response(response_data, self.get_refresh_token_name())
+
+    def _extract_token_expiry_date(self, response_data: Mapping[str, Any]) -> Any:
+        """
+        Extracts the token_expiry_date, like `expires_in` or `expires_at`, etc from the given response data.
+
+        Args:
+            response_data (Mapping[str, Any]): The response data from which to extract the token_expiry_date.
+
+        Returns:
+            str: The extracted token_expiry_date.
         """
+        return self._find_and_get_value_from_response(response_data, self.get_expires_in_name())
+
+    def _find_and_get_value_from_response(
+        self,
+        response_data: Mapping[str, Any],
+        key_name: str,
+        max_depth: int = 5,
+        current_depth: int = 0,
+    ) -> Any:
+        """
+        Recursively searches for a specified key in a nested dictionary or list and returns its value if found.
+
+        Args:
+            response_data (Mapping[str, Any]): The response data to search through, which can be a dictionary or a list.
+            key_name (str): The key to search for in the response data.
+            max_depth (int, optional): The maximum depth to search for the key to avoid infinite recursion. Defaults to 5.
+            current_depth (int, optional): The current depth of the recursion. Defaults to 0.
+
+        Returns:
+            Any: The value associated with the specified key if found, otherwise None.
+
+        Raises:
+            AirbyteTracedException: If the maximum recursion depth is reached without finding the key.
+        """
+        if current_depth > max_depth:
+            # this is needed to avoid an inf loop, possible with a very deep nesting observed.
+            message = f"The maximum level of recursion is reached. Couldn't find the speficied `{key_name}` in the response."
+            raise ResponseKeysMaxRecurtionReached(
+                internal_message=message, message=message, failure_type=FailureType.config_error
+            )
+
+        if isinstance(response_data, dict):
+            # get from the root level
+            if key_name in response_data:
+                return response_data[key_name]
+
+            # get from the nested object
+            for _, value in response_data.items():
+                result = self._find_and_get_value_from_response(
+                    value, key_name, max_depth, current_depth + 1
+                )
+                if result is not None:
+                    return result
+
+        # get from the nested array object
+        elif isinstance(response_data, list):
+            for item in response_data:
+                result = self._find_and_get_value_from_response(
+                    item, key_name, max_depth, current_depth + 1
+                )
+                if result is not None:
+                    return result
 
         return None
 
+    @property
+    def _message_repository(self) -> Optional[MessageRepository]:
+        """
+        The implementation can define a message_repository if it wants debugging logs for HTTP requests
+        """
+        return _NOOP_MESSAGE_REPOSITORY
+
+    def _log_response(self, response: requests.Response) -> None:
+        """
+        Logs the HTTP response using the message repository if it is available.
+
+        Args:
+            response (requests.Response): The HTTP response to log.
+        """
+        if self._message_repository:
+            self._message_repository.log_message(
+                Level.DEBUG,
+                lambda: format_http_message(
+                    response,
+                    "Refresh token",
+                    "Obtains access token",
+                    self._NO_STREAM_NAME,
+                    is_auxiliary=True,
+                ),
+            )
+
+    # ----------------
+    # ABSTR METHODS
+    # ----------------
+
     @abstractmethod
     def get_token_refresh_endpoint(self) -> Optional[str]:
         """Returns the endpoint to refresh the access token"""
@@ -295,23 +473,3 @@ class AbstractOauth2Authenticator(AuthBase):
     @abstractmethod
     def access_token(self, value: str) -> str:
         """Setter for the access token"""
-
-    @property
-    def _message_repository(self) -> Optional[MessageRepository]:
-        """
-        The implementation can define a message_repository if it wants debugging logs for HTTP requests
-        """
-        return _NOOP_MESSAGE_REPOSITORY
-
-    def _log_response(self, response: requests.Response) -> None:
-        if self._message_repository:
-            self._message_repository.log_message(
-                Level.DEBUG,
-                lambda: format_http_message(
-                    response,
-                    "Refresh token",
-                    "Obtains access token",
-                    self._NO_STREAM_NAME,
-                    is_auxiliary=True,
-                ),
-            )