airbyte_cdk.sources.declarative.partition_routers

20@dataclass
21class AsyncJobPartitionRouter(StreamSlicer):
22    """
23    Partition router that creates async jobs in a source API, periodically polls for job
24    completion, and supplies the completed job URL locations as stream slices so that
25    records can be extracted.
26    """
27
28    config: Config
29    parameters: InitVar[Mapping[str, Any]]
30    job_orchestrator_factory: Callable[[Iterable[StreamSlice]], AsyncJobOrchestrator]
31    stream_slicer: StreamSlicer = field(
32        default_factory=lambda: SinglePartitionRouter(parameters={})
33    )
34
35    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
36        self._job_orchestrator_factory = self.job_orchestrator_factory
37        self._job_orchestrator: Optional[AsyncJobOrchestrator] = None
38        self._parameters = parameters
39
40    def stream_slices(self) -> Iterable[StreamSlice]:
41        slices = self.stream_slicer.stream_slices()
42        self._job_orchestrator = self._job_orchestrator_factory(slices)
43
44        for completed_partition in self._job_orchestrator.create_and_get_completed_partitions():
45            yield StreamSlice(
46                partition=dict(completed_partition.stream_slice.partition),
47                cursor_slice=completed_partition.stream_slice.cursor_slice,
48                extra_fields={"jobs": list(completed_partition.jobs)},
49            )
50
51    def fetch_records(self, async_jobs: Iterable[AsyncJob]) -> Iterable[Mapping[str, Any]]:
52        """
53        This method of fetching records extends beyond what a PartitionRouter/StreamSlicer should
54        be responsible for. However, this was added in because the JobOrchestrator is required to
55        retrieve records. And without defining fetch_records() on this class, we're stuck with either
56        passing the JobOrchestrator to the AsyncRetriever or storing it on multiple classes.
57        """
58
59        if not self._job_orchestrator:
60            raise AirbyteTracedException(
61                message="Invalid state within AsyncJobRetriever. Please contact Airbyte Support",
62                internal_message="AsyncPartitionRepository is expected to be accessed only after `stream_slices`",
63                failure_type=FailureType.system_error,
64            )
65
66        return self._job_orchestrator.fetch_records(async_jobs=async_jobs)

 40@dataclass
 41class CartesianProductStreamSlicer(PartitionRouter):
 42    """
 43    Stream slicers that iterates over the cartesian product of input stream slicers
 44    Given 2 stream slicers with the following slices:
 45    A: [{"i": 0}, {"i": 1}, {"i": 2}]
 46    B: [{"s": "hello"}, {"s": "world"}]
 47    the resulting stream slices are
 48    [
 49        {"i": 0, "s": "hello"},
 50        {"i": 0, "s": "world"},
 51        {"i": 1, "s": "hello"},
 52        {"i": 1, "s": "world"},
 53        {"i": 2, "s": "hello"},
 54        {"i": 2, "s": "world"},
 55    ]
 56
 57    Attributes:
 58        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.
 59    """
 60
 61    stream_slicers: List[PartitionRouter]
 62    parameters: InitVar[Mapping[str, Any]]
 63
 64    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 65        check_for_substream_in_slicers(self.stream_slicers, self.logger.warning)
 66
 67    def get_request_params(
 68        self,
 69        *,
 70        stream_state: Optional[StreamState] = None,
 71        stream_slice: Optional[StreamSlice] = None,
 72        next_page_token: Optional[Mapping[str, Any]] = None,
 73    ) -> Mapping[str, Any]:
 74        return dict(
 75            ChainMap(
 76                *[  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
 77                    s.get_request_params(
 78                        stream_state=stream_state,
 79                        stream_slice=stream_slice,
 80                        next_page_token=next_page_token,
 81                    )
 82                    for s in self.stream_slicers
 83                ]
 84            )
 85        )
 86
 87    def get_request_headers(
 88        self,
 89        *,
 90        stream_state: Optional[StreamState] = None,
 91        stream_slice: Optional[StreamSlice] = None,
 92        next_page_token: Optional[Mapping[str, Any]] = None,
 93    ) -> Mapping[str, Any]:
 94        return dict(
 95            ChainMap(
 96                *[  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
 97                    s.get_request_headers(
 98                        stream_state=stream_state,
 99                        stream_slice=stream_slice,
100                        next_page_token=next_page_token,
101                    )
102                    for s in self.stream_slicers
103                ]
104            )
105        )
106
107    def get_request_body_data(
108        self,
109        *,
110        stream_state: Optional[StreamState] = None,
111        stream_slice: Optional[StreamSlice] = None,
112        next_page_token: Optional[Mapping[str, Any]] = None,
113    ) -> Mapping[str, Any]:
114        return dict(
115            ChainMap(
116                *[  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
117                    s.get_request_body_data(
118                        stream_state=stream_state,
119                        stream_slice=stream_slice,
120                        next_page_token=next_page_token,
121                    )
122                    for s in self.stream_slicers
123                ]
124            )
125        )
126
127    def get_request_body_json(
128        self,
129        *,
130        stream_state: Optional[StreamState] = None,
131        stream_slice: Optional[StreamSlice] = None,
132        next_page_token: Optional[Mapping[str, Any]] = None,
133    ) -> Mapping[str, Any]:
134        return dict(
135            ChainMap(
136                *[  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
137                    s.get_request_body_json(
138                        stream_state=stream_state,
139                        stream_slice=stream_slice,
140                        next_page_token=next_page_token,
141                    )
142                    for s in self.stream_slicers
143                ]
144            )
145        )
146
147    def stream_slices(self) -> Iterable[StreamSlice]:
148        sub_slices = (s.stream_slices() for s in self.stream_slicers)
149        product = itertools.product(*sub_slices)
150        for stream_slice_tuple in product:
151            partition = dict(ChainMap(*[s.partition for s in stream_slice_tuple]))  # type: ignore # ChainMap expects a MutableMapping[Never, Never] for reasons
152            cursor_slices = [s.cursor_slice for s in stream_slice_tuple if s.cursor_slice]
153            if len(cursor_slices) > 1:
154                raise ValueError(
155                    f"There should only be a single cursor slice. Found {cursor_slices}"
156                )
157            if cursor_slices:
158                cursor_slice = cursor_slices[0]
159            else:
160                cursor_slice = {}
161            yield StreamSlice(partition=partition, cursor_slice=cursor_slice)
162
163    def set_initial_state(self, stream_state: StreamState) -> None:
164        """
165        Parent stream states are not supported for cartesian product stream slicer
166        """
167        pass
168
169    def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
170        """
171        Parent stream states are not supported for cartesian product stream slicer
172        """
173        pass
174
175    @property
176    def logger(self) -> logging.Logger:
177        return logging.getLogger("airbyte.CartesianProductStreamSlicer")

 13@dataclass
 14class GroupingPartitionRouter(PartitionRouter):
 15    """
 16    A partition router that groups partitions from an underlying partition router into batches of a specified size.
 17    This is useful for APIs that support filtering by multiple partition keys in a single request.
 18
 19    Attributes:
 20        group_size (int): The number of partitions to include in each group.
 21        underlying_partition_router (PartitionRouter): The partition router whose output will be grouped.
 22        deduplicate (bool): If True, ensures unique partitions within each group by removing duplicates based on the partition key.
 23        config (Config): The connector configuration.
 24        parameters (Mapping[str, Any]): Additional parameters for interpolation and configuration.
 25    """
 26
 27    group_size: int
 28    underlying_partition_router: PartitionRouter
 29    config: Config
 30    deduplicate: bool = True
 31
 32    def __post_init__(self) -> None:
 33        self._state: Optional[Mapping[str, StreamState]] = {}
 34
 35    def stream_slices(self) -> Iterable[StreamSlice]:
 36        """
 37        Lazily groups partitions from the underlying partition router into batches of size `group_size`.
 38
 39        This method processes partitions one at a time from the underlying router, maintaining a batch buffer.
 40        When the buffer reaches `group_size` or the underlying router is exhausted, it yields a grouped slice.
 41        If deduplication is enabled, it tracks seen partition keys to ensure uniqueness within the current batch.
 42
 43        Yields:
 44            Iterable[StreamSlice]: An iterable of StreamSlice objects, where each slice contains a batch of partition values.
 45        """
 46        batch = []
 47        seen_keys = set()
 48
 49        # Iterate over partitions lazily from the underlying router
 50        for partition in self.underlying_partition_router.stream_slices():
 51            # Extract the partition key (assuming single key-value pair, e.g., {"board_ids": value})
 52            partition_keys = list(partition.partition.keys())
 53            # skip parent_slice as it is part of SubstreamPartitionRouter partition
 54            if "parent_slice" in partition_keys:
 55                partition_keys.remove("parent_slice")
 56            if len(partition_keys) != 1:
 57                raise ValueError(
 58                    f"GroupingPartitionRouter expects a single partition key-value pair. Got {partition.partition}"
 59                )
 60            key = partition.partition[partition_keys[0]]
 61
 62            # Skip duplicates if deduplication is enabled
 63            if self.deduplicate and key in seen_keys:
 64                continue
 65
 66            # Add partition to the batch
 67            batch.append(partition)
 68            if self.deduplicate:
 69                seen_keys.add(key)
 70
 71            # Yield the batch when it reaches the group_size
 72            if len(batch) == self.group_size:
 73                self._state = self.underlying_partition_router.get_stream_state()
 74                yield self._create_grouped_slice(batch)
 75                batch = []  # Reset the batch
 76
 77        self._state = self.underlying_partition_router.get_stream_state()
 78        # Yield any remaining partitions if the batch isn't empty
 79        if batch:
 80            yield self._create_grouped_slice(batch)
 81
 82    def _create_grouped_slice(self, batch: list[StreamSlice]) -> StreamSlice:
 83        """
 84        Creates a grouped StreamSlice from a batch of partitions, aggregating extra fields into a dictionary with list values.
 85
 86        Args:
 87            batch (list[StreamSlice]): A list of StreamSlice objects to group.
 88
 89        Returns:
 90            StreamSlice: A single StreamSlice with combined partition and extra field values.
 91        """
 92        # Combine partition values into a single dict with lists
 93        grouped_partition = {
 94            key: [p.partition.get(key) for p in batch] for key in batch[0].partition.keys()
 95        }
 96
 97        # Aggregate extra fields into a dict with list values
 98        extra_fields_dict = (
 99            {
100                key: [p.extra_fields.get(key) for p in batch]
101                for key in set().union(*(p.extra_fields.keys() for p in batch if p.extra_fields))
102            }
103            if any(p.extra_fields for p in batch)
104            else {}
105        )
106        return StreamSlice(
107            partition=grouped_partition,
108            cursor_slice={},  # Cursor is managed by the underlying router or incremental sync
109            extra_fields=extra_fields_dict,
110        )
111
112    def get_request_params(
113        self,
114        stream_state: Optional[StreamState] = None,
115        stream_slice: Optional[StreamSlice] = None,
116        next_page_token: Optional[Mapping[str, Any]] = None,
117    ) -> Mapping[str, Any]:
118        return {}
119
120    def get_request_headers(
121        self,
122        stream_state: Optional[StreamState] = None,
123        stream_slice: Optional[StreamSlice] = None,
124        next_page_token: Optional[Mapping[str, Any]] = None,
125    ) -> Mapping[str, Any]:
126        return {}
127
128    def get_request_body_data(
129        self,
130        stream_state: Optional[StreamState] = None,
131        stream_slice: Optional[StreamSlice] = None,
132        next_page_token: Optional[Mapping[str, Any]] = None,
133    ) -> Mapping[str, Any]:
134        return {}
135
136    def get_request_body_json(
137        self,
138        stream_state: Optional[StreamState] = None,
139        stream_slice: Optional[StreamSlice] = None,
140        next_page_token: Optional[Mapping[str, Any]] = None,
141    ) -> Mapping[str, Any]:
142        return {}
143
144    def set_initial_state(self, stream_state: StreamState) -> None:
145        """Delegate state initialization to the underlying partition router."""
146        self.underlying_partition_router.set_initial_state(stream_state)
147        self._state = self.underlying_partition_router.get_stream_state()
148
149    def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
150        """Delegate state retrieval to the underlying partition router."""
151        return self._state

 18@dataclass
 19class ListPartitionRouter(PartitionRouter):
 20    """
 21    Partition router that iterates over the values of a list
 22    If values is a string, then evaluate it as literal and assert the resulting literal is a list
 23
 24    Attributes:
 25        values (Union[str, List[str]]): The values to iterate over
 26        cursor_field (Union[InterpolatedString, str]): The name of the cursor field
 27        config (Config): The user-provided configuration as specified by the source's spec
 28        request_option (Optional[RequestOption]): The request option to configure the HTTP request
 29    """
 30
 31    values: Union[str, List[str]]
 32    cursor_field: Union[InterpolatedString, str]
 33    config: Config
 34    parameters: InitVar[Mapping[str, Any]]
 35    request_option: Optional[RequestOption] = None
 36
 37    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 38        if isinstance(self.values, str):
 39            self.values = InterpolatedString.create(self.values, parameters=parameters).eval(
 40                self.config
 41            )
 42        self._cursor_field = (
 43            InterpolatedString(string=self.cursor_field, parameters=parameters)
 44            if isinstance(self.cursor_field, str)
 45            else self.cursor_field
 46        )
 47
 48        self._cursor = None
 49
 50    def get_request_params(
 51        self,
 52        stream_state: Optional[StreamState] = None,
 53        stream_slice: Optional[StreamSlice] = None,
 54        next_page_token: Optional[Mapping[str, Any]] = None,
 55    ) -> Mapping[str, Any]:
 56        # Pass the stream_slice from the argument, not the cursor because the cursor is updated after processing the response
 57        return self._get_request_option(RequestOptionType.request_parameter, stream_slice)
 58
 59    def get_request_headers(
 60        self,
 61        stream_state: Optional[StreamState] = None,
 62        stream_slice: Optional[StreamSlice] = None,
 63        next_page_token: Optional[Mapping[str, Any]] = None,
 64    ) -> Mapping[str, Any]:
 65        # Pass the stream_slice from the argument, not the cursor because the cursor is updated after processing the response
 66        return self._get_request_option(RequestOptionType.header, stream_slice)
 67
 68    def get_request_body_data(
 69        self,
 70        stream_state: Optional[StreamState] = None,
 71        stream_slice: Optional[StreamSlice] = None,
 72        next_page_token: Optional[Mapping[str, Any]] = None,
 73    ) -> Mapping[str, Any]:
 74        # Pass the stream_slice from the argument, not the cursor because the cursor is updated after processing the response
 75        return self._get_request_option(RequestOptionType.body_data, stream_slice)
 76
 77    def get_request_body_json(
 78        self,
 79        stream_state: Optional[StreamState] = None,
 80        stream_slice: Optional[StreamSlice] = None,
 81        next_page_token: Optional[Mapping[str, Any]] = None,
 82    ) -> Mapping[str, Any]:
 83        # Pass the stream_slice from the argument, not the cursor because the cursor is updated after processing the response
 84        return self._get_request_option(RequestOptionType.body_json, stream_slice)
 85
 86    def stream_slices(self) -> Iterable[StreamSlice]:
 87        return [
 88            StreamSlice(
 89                partition={self._cursor_field.eval(self.config): slice_value}, cursor_slice={}
 90            )
 91            for slice_value in self.values
 92        ]
 93
 94    def _get_request_option(
 95        self, request_option_type: RequestOptionType, stream_slice: Optional[StreamSlice]
 96    ) -> Mapping[str, Any]:
 97        if (
 98            self.request_option
 99            and self.request_option.inject_into == request_option_type
100            and stream_slice
101        ):
102            slice_value = stream_slice.get(self._cursor_field.eval(self.config))
103            if slice_value:
104                options: MutableMapping[str, Any] = {}
105                self.request_option.inject_into_request(options, slice_value, self.config)
106                return options
107            else:
108                return {}
109        else:
110            return {}
111
112    def set_initial_state(self, stream_state: StreamState) -> None:
113        """
114        ListPartitionRouter doesn't have parent streams
115        """
116        pass
117
118    def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
119        """
120        ListPartitionRouter doesn't have parent streams
121        """
122        pass

13@dataclass
14class SinglePartitionRouter(PartitionRouter):
15    """Partition router returning only a stream slice"""
16
17    parameters: InitVar[Mapping[str, Any]]
18
19    def get_request_params(
20        self,
21        stream_state: Optional[StreamState] = None,
22        stream_slice: Optional[StreamSlice] = None,
23        next_page_token: Optional[Mapping[str, Any]] = None,
24    ) -> Mapping[str, Any]:
25        return {}
26
27    def get_request_headers(
28        self,
29        stream_state: Optional[StreamState] = None,
30        stream_slice: Optional[StreamSlice] = None,
31        next_page_token: Optional[Mapping[str, Any]] = None,
32    ) -> Mapping[str, Any]:
33        return {}
34
35    def get_request_body_data(
36        self,
37        stream_state: Optional[StreamState] = None,
38        stream_slice: Optional[StreamSlice] = None,
39        next_page_token: Optional[Mapping[str, Any]] = None,
40    ) -> Mapping[str, Any]:
41        return {}
42
43    def get_request_body_json(
44        self,
45        stream_state: Optional[StreamState] = None,
46        stream_slice: Optional[StreamSlice] = None,
47        next_page_token: Optional[Mapping[str, Any]] = None,
48    ) -> Mapping[str, Any]:
49        return {}
50
51    def stream_slices(self) -> Iterable[StreamSlice]:
52        yield StreamSlice(partition={}, cursor_slice={})
53
54    def set_initial_state(self, stream_state: StreamState) -> None:
55        """
56        SinglePartitionRouter doesn't have parent streams
57        """
58        pass
59
60    def get_stream_state(self) -> Optional[Mapping[str, StreamState]]:
61        """
62        SinglePartitionRouter doesn't have parent streams
63        """
64        pass