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

airbyte_cdk/sources/declarative/partition_routers/__init__.py

@@ -0,0 +1,29 @@
+#
+# Copyright (c) 2022 Airbyte, Inc., all rights reserved.
+#
+
+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__ = [
+    "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 import AsyncJob
+from airbyte_cdk.sources.declarative.async_job.job_orchestrator import (
+    AsyncJobOrchestrator,
+)
+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),
+                cursor_slice=completed_partition.stream_slice.cursor_slice,
+                extra_fields={"jobs": list(completed_partition.jobs)},
+            )
+
+    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
+        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(async_jobs=async_jobs)

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

@@ -0,0 +1,121 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from dataclasses import InitVar, dataclass
+from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Union
+
+from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
+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(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
+
+    Attributes:
+        values (Union[str, List[str]]): The values to iterate over
+        cursor_field (Union[InterpolatedString, str]): The name of the cursor field
+        config (Config): The user-provided configuration as specified by the source's spec
+        request_option (Optional[RequestOption]): The request option to configure the HTTP request
+    """
+
+    values: Union[str, List[str]]
+    cursor_field: Union[InterpolatedString, str]
+    config: Config
+    parameters: InitVar[Mapping[str, Any]]
+    request_option: Optional[RequestOption] = None
+
+    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._cursor_field = (
+            InterpolatedString(string=self.cursor_field, parameters=parameters)
+            if isinstance(self.cursor_field, str)
+            else self.cursor_field
+        )
+
+        self._cursor = None
+
+    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]:
+        # Pass the stream_slice from the argument, not the cursor because the cursor is updated after processing the response
+        return self._get_request_option(RequestOptionType.request_parameter, stream_slice)
+
+    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]:
+        # Pass the stream_slice from the argument, not the cursor because the cursor is updated after processing the response
+        return self._get_request_option(RequestOptionType.header, stream_slice)
+
+    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]:
+        # Pass the stream_slice from the argument, not the cursor because the cursor is updated after processing the response
+        return self._get_request_option(RequestOptionType.body_data, stream_slice)
+
+    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]:
+        # Pass the stream_slice from the argument, not the cursor because the cursor is updated after processing the response
+        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
+        ]
+
+    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:
+                options: MutableMapping[str, Any] = {}
+                self.request_option.inject_into_request(options, slice_value, self.config)
+                return options
+            else:
+                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

@@ -0,0 +1,63 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from dataclasses import InitVar, dataclass
+from typing import Any, Iterable, Mapping, Optional
+
+from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
+from airbyte_cdk.sources.types import StreamSlice, StreamState
+
+
+@dataclass
+class SinglePartitionRouter(PartitionRouter):
+    """Partition router returning only a stream slice"""
+
+    parameters: InitVar[Mapping[str, Any]]
+
+    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 {}
+
+    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 {}
+
+    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 {}
+
+    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 {}
+
+    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