airbyte-cdk 0.72.1__py3-none-any.whl → 6.17.1.dev1__py3-none-any.whl

airbyte_cdk/sources/declarative/partition_routers/__init__.py

@@ -2,8 +2,28 @@
 # Copyright (c) 2022 Airbyte, Inc., all rights reserved.
 #
 
-from airbyte_cdk.sources.declarative.partition_routers.list_partition_router import ListPartitionRouter
-from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import SinglePartitionRouter
-from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import SubstreamPartitionRouter
+from airbyte_cdk.sources.declarative.partition_routers.async_job_partition_router import (
+    AsyncJobPartitionRouter,
+)
+from airbyte_cdk.sources.declarative.partition_routers.cartesian_product_stream_slicer import (
+    CartesianProductStreamSlicer,
+)
+from airbyte_cdk.sources.declarative.partition_routers.list_partition_router import (
+    ListPartitionRouter,
+)
+from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
+from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import (
+    SinglePartitionRouter,
+)
+from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
+    SubstreamPartitionRouter,
+)
 
-__all__ = ["ListPartitionRouter", "SinglePartitionRouter", "SubstreamPartitionRouter"]
+__all__ = [
+    "AsyncJobPartitionRouter",
+    "CartesianProductStreamSlicer",
+    "ListPartitionRouter",
+    "SinglePartitionRouter",
+    "SubstreamPartitionRouter",
+    "PartitionRouter",
+]

airbyte_cdk/sources/declarative/partition_routers/async_job_partition_router.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+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_orchestrator import (
+    AsyncJobOrchestrator,
+    AsyncPartition,
+)
+from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import (
+    SinglePartitionRouter,
+)
+from airbyte_cdk.sources.streams.concurrent.partitions.stream_slicer import StreamSlicer
+from airbyte_cdk.sources.types import Config, StreamSlice
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+
+
+@dataclass
+class AsyncJobPartitionRouter(StreamSlicer):
+    """
+    Partition router that creates async jobs in a source API, periodically polls for job
+    completion, and supplies the completed job URL locations as stream slices so that
+    records can be extracted.
+    """
+
+    config: Config
+    parameters: InitVar[Mapping[str, Any]]
+    job_orchestrator_factory: Callable[[Iterable[StreamSlice]], AsyncJobOrchestrator]
+    stream_slicer: StreamSlicer = field(
+        default_factory=lambda: SinglePartitionRouter(parameters={})
+    )
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        self._job_orchestrator_factory = self.job_orchestrator_factory
+        self._job_orchestrator: Optional[AsyncJobOrchestrator] = None
+        self._parameters = parameters
+
+    def stream_slices(self) -> Iterable[StreamSlice]:
+        slices = self.stream_slicer.stream_slices()
+        self._job_orchestrator = self._job_orchestrator_factory(slices)
+
+        for completed_partition in self._job_orchestrator.create_and_get_completed_partitions():
+            yield StreamSlice(
+                partition=dict(completed_partition.stream_slice.partition)
+                | {"partition": completed_partition},
+                cursor_slice=completed_partition.stream_slice.cursor_slice,
+            )
+
+    def fetch_records(self, partition: AsyncPartition) -> 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
+        retrieve records. And without defining fetch_records() on this class, we're stuck with either
+        passing the JobOrchestrator to the AsyncRetriever or storing it on multiple classes.
+        """
+
+        if not self._job_orchestrator:
+            raise AirbyteTracedException(
+                message="Invalid state within AsyncJobRetriever. Please contact Airbyte Support",
+                internal_message="AsyncPartitionRepository is expected to be accessed only after `stream_slices`",
+                failure_type=FailureType.system_error,
+            )
+
+        return self._job_orchestrator.fetch_records(partition=partition)

airbyte_cdk/sources/declarative/partition_routers/cartesian_product_stream_slicer.py

@@ -0,0 +1,176 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import itertools
+import logging
+from collections import ChainMap
+from collections.abc import Callable
+from dataclasses import InitVar, dataclass
+from typing import Any, Iterable, List, Mapping, Optional
+
+from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
+from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
+    SubstreamPartitionRouter,
+)
+from airbyte_cdk.sources.types import StreamSlice, StreamState
+
+
+def check_for_substream_in_slicers(
+    slicers: Iterable[PartitionRouter], log_warning: Callable[[str], None]
+) -> None:
+    """
+    Recursively checks for the presence of SubstreamPartitionRouter within slicers.
+    Logs a warning if a SubstreamPartitionRouter is found within a CartesianProductStreamSlicer.
+
+    Args:
+        slicers (Iterable[PartitionRouter]): The list of slicers to check.
+        log_warning (Callable): Logging function to record warnings.
+    """
+    for slicer in slicers:
+        if isinstance(slicer, SubstreamPartitionRouter):
+            log_warning("Parent state handling is not supported for CartesianProductStreamSlicer.")
+            return
+        elif isinstance(slicer, CartesianProductStreamSlicer):
+            # Recursively check sub-slicers within CartesianProductStreamSlicer
+            check_for_substream_in_slicers(slicer.stream_slicers, log_warning)
+
+
+@dataclass
+class CartesianProductStreamSlicer(PartitionRouter):
+    """
+    Stream slicers that iterates over the cartesian product of input stream slicers
+    Given 2 stream slicers with the following slices:
+    A: [{"i": 0}, {"i": 1}, {"i": 2}]
+    B: [{"s": "hello"}, {"s": "world"}]
+    the resulting stream slices are
+    [
+        {"i": 0, "s": "hello"},
+        {"i": 0, "s": "world"},
+        {"i": 1, "s": "hello"},
+        {"i": 1, "s": "world"},
+        {"i": 2, "s": "hello"},
+        {"i": 2, "s": "world"},
+    ]
+
+    Attributes:
+        stream_slicers (List[PartitionRouter]): Underlying stream slicers. The RequestOptions (e.g: Request headers, parameters, etc..) returned by this slicer are the combination of the RequestOptions of its input slicers. If there are conflicts e.g: two slicers define the same header or request param, the conflict is resolved by taking the value from the first slicer, where ordering is determined by the order in which slicers were input to this composite slicer.
+    """
+
+    stream_slicers: List[PartitionRouter]
+    parameters: InitVar[Mapping[str, Any]]
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        check_for_substream_in_slicers(self.stream_slicers, self.logger.warning)
+
+    def get_request_params(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return dict(
+            ChainMap(
+                *[  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
+                    s.get_request_params(
+                        stream_state=stream_state,
+                        stream_slice=stream_slice,
+                        next_page_token=next_page_token,
+                    )
+                    for s in self.stream_slicers
+                ]
+            )
+        )
+
+    def get_request_headers(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return dict(
+            ChainMap(
+                *[  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
+                    s.get_request_headers(
+                        stream_state=stream_state,
+                        stream_slice=stream_slice,
+                        next_page_token=next_page_token,
+                    )
+                    for s in self.stream_slicers
+                ]
+            )
+        )
+
+    def get_request_body_data(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return dict(
+            ChainMap(
+                *[  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
+                    s.get_request_body_data(
+                        stream_state=stream_state,
+                        stream_slice=stream_slice,
+                        next_page_token=next_page_token,
+                    )
+                    for s in self.stream_slicers
+                ]
+            )
+        )
+
+    def get_request_body_json(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return dict(
+            ChainMap(
+                *[  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
+                    s.get_request_body_json(
+                        stream_state=stream_state,
+                        stream_slice=stream_slice,
+                        next_page_token=next_page_token,
+                    )
+                    for s in self.stream_slicers
+                ]
+            )
+        )
+
+    def stream_slices(self) -> Iterable[StreamSlice]:
+        sub_slices = (s.stream_slices() for s in self.stream_slicers)
+        product = itertools.product(*sub_slices)
+        for stream_slice_tuple in product:
+            partition = dict(ChainMap(*[s.partition for s in stream_slice_tuple]))  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
+            cursor_slices = [s.cursor_slice for s in stream_slice_tuple if s.cursor_slice]
+            if len(cursor_slices) > 1:
+                raise ValueError(
+                    f"There should only be a single cursor slice. Found {cursor_slices}"
+                )
+            if cursor_slices:
+                cursor_slice = cursor_slices[0]
+            else:
+                cursor_slice = {}
+            yield StreamSlice(partition=partition, cursor_slice=cursor_slice)
+
+    def set_initial_state(self, stream_state: StreamState) -> None:
+        """
+        Parent stream states are not supported for cartesian product stream slicer
+        """
+        pass
+
+    def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
+        """
+        Parent stream states are not supported for cartesian product stream slicer
+        """
+        pass
+
+    @property
+    def logger(self) -> logging.Logger:
+        return logging.getLogger("airbyte.CartesianProductStreamSlicer")

airbyte_cdk/sources/declarative/partition_routers/list_partition_router.py

@@ -6,13 +6,16 @@ from dataclasses import InitVar, dataclass
 from typing import Any, Iterable, List, Mapping, Optional, Union
 
 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.stream_slicers.stream_slicer import StreamSlicer
-from airbyte_cdk.sources.declarative.types import Config, StreamSlice, StreamState
+from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
+from airbyte_cdk.sources.declarative.requesters.request_option import (
+    RequestOption,
+    RequestOptionType,
+)
+from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
 
 
 @dataclass
-class ListPartitionRouter(StreamSlicer):
+class ListPartitionRouter(PartitionRouter):
     """
     Partition router that iterates over the values of a list
     If values is a string, then evaluate it as literal and assert the resulting literal is a list
@@ -32,9 +35,13 @@ class ListPartitionRouter(StreamSlicer):
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         if isinstance(self.values, str):
-            self.values = InterpolatedString.create(self.values, parameters=parameters).eval(self.config)
+            self.values = InterpolatedString.create(self.values, parameters=parameters).eval(
+                self.config
+            )
         self._cursor_field = (
-            InterpolatedString(string=self.cursor_field, parameters=parameters) if isinstance(self.cursor_field, str) else self.cursor_field
+            InterpolatedString(string=self.cursor_field, parameters=parameters)
+            if isinstance(self.cursor_field, str)
+            else self.cursor_field
         )
 
         self._cursor = None
@@ -76,10 +83,21 @@ class ListPartitionRouter(StreamSlicer):
         return self._get_request_option(RequestOptionType.body_json, stream_slice)
 
     def stream_slices(self) -> Iterable[StreamSlice]:
-        return [StreamSlice(partition={self._cursor_field.eval(self.config): slice_value}, cursor_slice={}) for slice_value in self.values]
+        return [
+            StreamSlice(
+                partition={self._cursor_field.eval(self.config): slice_value}, cursor_slice={}
+            )
+            for slice_value in self.values
+        ]
 
-    def _get_request_option(self, request_option_type: RequestOptionType, stream_slice: Optional[StreamSlice]) -> Mapping[str, Any]:
-        if self.request_option and self.request_option.inject_into == request_option_type and stream_slice:
+    def _get_request_option(
+        self, request_option_type: RequestOptionType, stream_slice: Optional[StreamSlice]
+    ) -> Mapping[str, Any]:
+        if (
+            self.request_option
+            and self.request_option.inject_into == request_option_type
+            and stream_slice
+        ):
             slice_value = stream_slice.get(self._cursor_field.eval(self.config))
             if slice_value:
                 return {self.request_option.field_name.eval(self.config): slice_value}  # type: ignore # field_name is always casted to InterpolatedString
@@ -87,3 +105,15 @@ class ListPartitionRouter(StreamSlicer):
                 return {}
         else:
             return {}
+
+    def set_initial_state(self, stream_state: StreamState) -> None:
+        """
+        ListPartitionRouter doesn't have parent streams
+        """
+        pass
+
+    def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
+        """
+        ListPartitionRouter doesn't have parent streams
+        """
+        pass

airbyte_cdk/sources/declarative/partition_routers/partition_router.py

@@ -0,0 +1,62 @@
+#
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+#
+
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Mapping, Optional
+
+from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer
+from airbyte_cdk.sources.types import StreamState
+
+
+@dataclass
+class PartitionRouter(StreamSlicer):
+    """
+    Base class for partition routers.
+    Methods:
+        set_parent_state(stream_state): Set the state of the parent streams.
+        get_parent_state(): Get the state of the parent streams.
+    """
+
+    @abstractmethod
+    def set_initial_state(self, stream_state: StreamState) -> None:
+        """
+        Set the state of the parent streams.
+
+        This method should only be implemented if the slicer is based on some parent stream and needs to read this stream
+        incrementally using the state.
+
+        Args:
+            stream_state (StreamState): The state of the streams to be set. The expected format is a dictionary that includes
+                                        'parent_state' which is a dictionary of parent state names to their corresponding state.
+                Example:
+                {
+                    "parent_state": {
+                        "parent_stream_name_1": { ... },
+                        "parent_stream_name_2": { ... },
+                        ...
+                    }
+                }
+        """
+
+    @abstractmethod
+    def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
+        """
+        Get the state of the parent streams.
+
+        This method should only be implemented if the slicer is based on some parent stream and needs to read this stream
+        incrementally using the state.
+
+        Returns:
+            Optional[Mapping[str, StreamState]]: The current state of the parent streams in a dictionary format.
+                 The returned format will be:
+                 {
+                     "parent_stream_name1": {
+                         "last_updated": "2023-05-27T00:00:00Z"
+                     },
+                     "parent_stream_name2": {
+                         "last_updated": "2023-05-27T00:00:00Z"
+                     }
+                 }
+        """

airbyte_cdk/sources/declarative/partition_routers/single_partition_router.py

@@ -5,12 +5,12 @@
 from dataclasses import InitVar, dataclass
 from typing import Any, Iterable, Mapping, Optional
 
-from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer
-from airbyte_cdk.sources.declarative.types import StreamSlice, StreamState
+from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
+from airbyte_cdk.sources.types import StreamSlice, StreamState
 
 
 @dataclass
-class SinglePartitionRouter(StreamSlicer):
+class SinglePartitionRouter(PartitionRouter):
     """Partition router returning only a stream slice"""
 
     parameters: InitVar[Mapping[str, Any]]
@@ -49,3 +49,15 @@ class SinglePartitionRouter(StreamSlicer):
 
     def stream_slices(self) -> Iterable[StreamSlice]:
         yield StreamSlice(partition={}, cursor_slice={})
+
+    def set_initial_state(self, stream_state: StreamState) -> None:
+        """
+        SinglePartitionRouter doesn't have parent streams
+        """
+        pass
+
+    def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
+        """
+        SinglePartitionRouter doesn't have parent streams
+        """
+        pass