airbyte_cdk.sources.declarative.retrievers

15class Retriever:
16    """
17    Responsible for fetching a stream's records from an HTTP API source.
18
19    A stream creates one retriever and reads all its partitions through it: `read_records` is called from several
20    threads at once, one call per stream slice. Implementations, custom ones included, must keep per-read state such
21    as page tokens, counters or recursion depth in local variables rather than on the instance or its components.
22    """
23
24    @abstractmethod
25    def read_records(
26        self,
27        records_schema: Mapping[str, Any],
28        stream_slice: Optional[StreamSlice] = None,
29    ) -> Iterable[StreamData]:
30        """
31        Fetch a stream's records from an HTTP API source
32
33        :param records_schema: json schema to describe record
34        :param stream_slice: The stream slice to read data for
35        :return: The records read from the API source
36        """
37
38    @deprecated("Stream slicing is being moved to the stream level.")
39    def stream_slices(self) -> Iterable[Optional[StreamSlice]]:
40        """Does nothing as this method is deprecated, so underlying Retriever implementations
41        do not need to implement this.
42        """
43        yield from []
44
45    @property
46    @deprecated("State management is being moved to the stream level.")
47    def state(self) -> StreamState:
48        """
49        Does nothing as this method is deprecated, so underlying Retriever implementations
50        do not need to implement this.
51        """
52        return {}
53
54    @state.setter
55    @deprecated("State management is being moved to the stream level.")
56    def state(self, value: StreamState) -> None:
57        """
58        Does nothing as this method is deprecated, so underlying Retriever implementations
59        do not need to implement this.
60        """
61        pass

 93@dataclass
 94class SimpleRetriever(Retriever):
 95    """
 96    Retrieves records by synchronously sending requests to fetch records.
 97
 98    The retriever acts as an orchestrator between the requester, the record selector, the paginator, and the stream slicer.
 99
100    For each stream slice, submit requests until there are no more pages of records to fetch.
101
102    This retriever currently inherits from HttpStream to reuse the request submission and pagination machinery.
103    As a result, some of the parameters passed to some methods are unused.
104    The two will be decoupled in a future release.
105
106    Attributes:
107        stream_name (str): The stream's name
108        stream_primary_key (Optional[Union[str, List[str], List[List[str]]]]): The stream's primary key
109        requester (Requester): The HTTP requester
110        record_selector (HttpSelector): The record selector
111        paginator (Optional[Paginator]): The paginator
112        stream_slicer (Optional[StreamSlicer]): The stream slicer
113        parameters (Mapping[str, Any]): Additional runtime parameters to be used for string interpolation
114        post_pagination_filter (Optional[ClientSideIncrementalRecordFilterDecorator]): Set for data feed streams only.
115            Records the cursor considers already synced are dropped once pagination has observed them
116        page_size_reduction (Optional[PageSizeReduction]): How much to shrink the page size when an error handler
117            resolves to `ResponseAction.REDUCE_PAGE_SIZE`. `None` disables page size reduction entirely.
118            It is immutable configuration; the page size in effect lives in a `PageSizeReducer` that
119            `_read_pages` creates per call, so the retriever and its paginator - both shared by every
120            partition of the stream, read concurrently - stay stateless. When `page_size_reduction` is
121            `None` no reducer is created and `_read_pages` keeps its previous behaviour
122        request_window_splitting (Optional[RequestWindowSplitting]): Policy applied when an error handler
123            resolves to `ResponseAction.SPLIT_REQUEST_WINDOW`, or when `RequestWindowSplitRequiredException`
124            is raised directly by custom code. `None` disables request-window splitting entirely: `read_records`
125            converts the exception into `RequestWindowSplitNotSupportedException` instead of attempting a
126            split it cannot perform.
127        request_window_splitter: Bound method - normally the stream's own cursor's `split_request_window` -
128            asked to replace a failing `StreamSlice` with smaller children. `None` has the same effect as
129            `request_window_splitting` being `None`.
130        page_observer (Optional[Callable[[PageCompleted], None]]): Observes completed response pages; set
131            only on streams using `SyncTokenCursor` and `None` otherwise.
132    """
133
134    requester: Requester
135    record_selector: HttpSelector
136    config: Config
137    parameters: InitVar[Mapping[str, Any]]
138    name: str
139    _name: Union[InterpolatedString, str] = field(init=False, repr=False, default="")
140    primary_key: Optional[Union[str, List[str], List[List[str]]]]
141    _primary_key: str = field(init=False, repr=False, default="")
142    paginator: Optional[Paginator] = None
143    stream_slicer: StreamSlicer = field(
144        default_factory=lambda: SinglePartitionRouter(parameters={})
145    )
146    request_option_provider: RequestOptionsProvider = field(
147        default_factory=lambda: DefaultRequestOptionsProvider(parameters={})
148    )
149    ignore_stream_slicer_parameters_on_paginated_requests: bool = False
150    additional_query_properties: Optional[QueryProperties] = None
151    log_formatter: Optional[Callable[[requests.Response], Any]] = None
152    pagination_tracker_factory: Callable[[], PaginationTracker] = field(
153        default_factory=lambda: lambda: PaginationTracker()
154    )
155    post_pagination_filter: Optional[ClientSideIncrementalRecordFilterDecorator] = None
156    page_size_reduction: Optional[PageSizeReduction] = None
157    request_window_splitting: Optional[RequestWindowSplitting] = None
158    request_window_splitter: Optional[
159        Callable[[StreamSlice, Optional[datetime.timedelta]], Optional[List[StreamSlice]]]
160    ] = None
161    page_observer: Optional[Callable[["PageCompleted"], None]] = None
162
163    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
164        self._paginator = self.paginator or NoPagination(parameters=parameters)
165        self._parameters = parameters
166        self._name = (
167            InterpolatedString(self._name, parameters=parameters)
168            if isinstance(self._name, str)
169            else self._name
170        )
171
172    @property  # type: ignore
173    def name(self) -> str:
174        """
175        :return: Stream name
176        """
177        return (
178            str(self._name.eval(self.config))
179            if isinstance(self._name, InterpolatedString)
180            else self._name
181        )
182
183    @name.setter
184    def name(self, value: str) -> None:
185        if not isinstance(value, property):
186            self._name = value
187
188    def _get_mapping(
189        self, method: Callable[..., Optional[Union[Mapping[str, Any], str]]], **kwargs: Any
190    ) -> Tuple[Union[Mapping[str, Any], str], Set[str]]:
191        """
192        Get mapping from the provided method, and get the keys of the mapping.
193        If the method returns a string, it will return the string and an empty set.
194        If the method returns a dict, it will return the dict and its keys.
195        """
196        mapping = method(**kwargs) or {}
197        keys = set(mapping.keys()) if not isinstance(mapping, str) else set()
198        return mapping, keys
199
200    def _get_request_options(
201        self,
202        stream_slice: Optional[StreamSlice],
203        next_page_token: Optional[Mapping[str, Any]],
204        paginator_method: Callable[..., Optional[Union[Mapping[str, Any], str]]],
205        stream_slicer_method: Callable[..., Optional[Union[Mapping[str, Any], str]]],
206        page_size_override: Optional[int] = None,
207    ) -> Union[Mapping[str, Any], str]:
208        """
209        Get the request_option from the paginator and the stream slicer.
210        Raise a ValueError if there's a key collision
211        Returned merged mapping otherwise
212        """
213        is_body_json = paginator_method.__name__ == "get_request_body_json"
214
215        mappings = [
216            paginator_method(
217                stream_slice=stream_slice,
218                next_page_token=next_page_token,
219                **page_size_override_kwargs(page_size_override),
220            ),
221        ]
222        if not next_page_token or not self.ignore_stream_slicer_parameters_on_paginated_requests:
223            mappings.append(
224                stream_slicer_method(
225                    stream_slice=stream_slice,
226                    next_page_token=next_page_token,
227                )
228            )
229        return combine_mappings(mappings, allow_same_value_merge=is_body_json)
230
231    def _request_headers(
232        self,
233        stream_slice: Optional[StreamSlice] = None,
234        next_page_token: Optional[Mapping[str, Any]] = None,
235        page_size_override: Optional[int] = None,
236    ) -> Mapping[str, Any]:
237        """
238        Specifies request headers.
239        Authentication headers will overwrite any overlapping headers returned from this method.
240        """
241        headers = self._get_request_options(
242            stream_slice,
243            next_page_token,
244            self._paginator.get_request_headers,
245            self.request_option_provider.get_request_headers,
246            **page_size_override_kwargs(page_size_override),
247        )
248        if isinstance(headers, str):
249            raise ValueError("Request headers cannot be a string")
250        return {str(k): str(v) for k, v in headers.items()}
251
252    def _request_params(
253        self,
254        stream_slice: Optional[StreamSlice] = None,
255        next_page_token: Optional[Mapping[str, Any]] = None,
256        page_size_override: Optional[int] = None,
257    ) -> Mapping[str, Any]:
258        """
259        Specifies the query parameters that should be set on an outgoing HTTP request given the inputs.
260
261        E.g: you might want to define query parameters for paging if next_page_token is not None.
262        """
263        params = self._get_request_options(
264            stream_slice,
265            next_page_token,
266            self._paginator.get_request_params,
267            self.request_option_provider.get_request_params,
268            **page_size_override_kwargs(page_size_override),
269        )
270        if isinstance(params, str):
271            raise ValueError("Request params cannot be a string")
272        return params
273
274    def _request_body_data(
275        self,
276        stream_slice: Optional[StreamSlice] = None,
277        next_page_token: Optional[Mapping[str, Any]] = None,
278        page_size_override: Optional[int] = None,
279    ) -> Union[Mapping[str, Any], str]:
280        """
281        Specifies how to populate the body of the request with a non-JSON payload.
282
283        If returns a ready text that it will be sent as is.
284        If returns a dict that it will be converted to a urlencoded form.
285        E.g. {"key1": "value1", "key2": "value2"} => "key1=value1&key2=value2"
286
287        At the same time only one of the 'request_body_data' and 'request_body_json' functions can be overridden.
288        """
289        return self._get_request_options(
290            stream_slice,
291            next_page_token,
292            self._paginator.get_request_body_data,
293            self.request_option_provider.get_request_body_data,
294            **page_size_override_kwargs(page_size_override),
295        )
296
297    def _request_body_json(
298        self,
299        stream_slice: Optional[StreamSlice] = None,
300        next_page_token: Optional[Mapping[str, Any]] = None,
301        page_size_override: Optional[int] = None,
302    ) -> Optional[Mapping[str, Any]]:
303        """
304        Specifies how to populate the body of the request with a JSON payload.
305
306        At the same time only one of the 'request_body_data' and 'request_body_json' functions can be overridden.
307        """
308        body_json = self._get_request_options(
309            stream_slice,
310            next_page_token,
311            self._paginator.get_request_body_json,
312            self.request_option_provider.get_request_body_json,
313            **page_size_override_kwargs(page_size_override),
314        )
315        if isinstance(body_json, str):
316            raise ValueError("Request body json cannot be a string")
317        return body_json
318
319    def _paginator_path(
320        self,
321        next_page_token: Optional[Mapping[str, Any]] = None,
322        stream_slice: Optional[StreamSlice] = None,
323    ) -> Optional[str]:
324        """
325        If the paginator points to a path, follow it, else return nothing so the requester is used.
326        :param next_page_token:
327        :return:
328        """
329        return self._paginator.path(
330            next_page_token=next_page_token,
331            stream_state={},  # stream_state as an interpolation context is deprecated
332            stream_slice=stream_slice,
333        )
334
335    def _parse_response(
336        self,
337        response: Optional[requests.Response],
338        records_schema: Mapping[str, Any],
339        stream_slice: Optional[StreamSlice] = None,
340        next_page_token: Optional[Mapping[str, Any]] = None,
341    ) -> Iterable[Record]:
342        if not response:
343            yield from []
344        else:
345            yield from self.record_selector.select_records(
346                response=response,
347                stream_state={},  # stream_state as an interpolation context is deprecated
348                records_schema=records_schema,
349                stream_slice=stream_slice,
350                next_page_token=next_page_token,
351            )
352
353    @property  # type: ignore
354    def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
355        """The stream's primary key"""
356        return self._primary_key
357
358    @primary_key.setter
359    def primary_key(self, value: str) -> None:
360        if not isinstance(value, property):
361            self._primary_key = value
362
363    def _next_page_token(
364        self,
365        response: requests.Response,
366        last_page_size: int,
367        last_record: Optional[Record],
368        last_page_token_value: Optional[Any],
369        page_size_override: Optional[int] = None,
370        stream_slice: Optional[StreamSlice] = None,
371    ) -> Optional[Mapping[str, Any]]:
372        """
373        Specifies a pagination strategy.
374
375        The value returned from this method is passed to most other methods in this class. Use it to form a request e.g: set headers or query params.
376
377        :return: The token for the next page from the input response object. Returning None means there are no more pages to read in this response.
378        """
379        return self._paginator.next_page_token(
380            response=response,
381            last_page_size=last_page_size,
382            last_record=last_record,
383            last_page_token_value=last_page_token_value,
384            **page_size_override_kwargs(page_size_override),
385            **stream_slice_kwargs(self._paginator.next_page_token, stream_slice),
386        )
387
388    def _fetch_next_page(
389        self,
390        stream_slice: StreamSlice,
391        next_page_token: Optional[Mapping[str, Any]] = None,
392        page_size_override: Optional[int] = None,
393    ) -> Optional[requests.Response]:
394        return self.requester.send_request(
395            path=self._paginator_path(
396                next_page_token=next_page_token,
397                stream_slice=stream_slice,
398            ),
399            stream_state={},  # stream_state as an interpolation context is deprecated
400            stream_slice=stream_slice,
401            next_page_token=next_page_token,
402            request_headers=self._request_headers(
403                stream_slice=stream_slice,
404                next_page_token=next_page_token,
405                **page_size_override_kwargs(page_size_override),
406            ),
407            request_params=self._request_params(
408                stream_slice=stream_slice,
409                next_page_token=next_page_token,
410                **page_size_override_kwargs(page_size_override),
411            ),
412            request_body_data=self._request_body_data(
413                stream_slice=stream_slice,
414                next_page_token=next_page_token,
415                **page_size_override_kwargs(page_size_override),
416            ),
417            request_body_json=self._request_body_json(
418                stream_slice=stream_slice,
419                next_page_token=next_page_token,
420                **page_size_override_kwargs(page_size_override),
421            ),
422            log_formatter=self.log_formatter,
423        )
424
425    # This logic is similar to _read_pages in the HttpStream class. When making changes here, consider making changes there as well.
426    def _read_pages(
427        self,
428        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
429        stream_slice: StreamSlice,
430    ) -> Iterable[Record]:
431        original_stream_slice = stream_slice
432        pagination_tracker = self.pagination_tracker_factory()
433        page_size_reducer = (
434            PageSizeReducer(
435                self.page_size_reduction,
436                self._paginator.get_page_size(),
437                stream_name=self.name,
438            )
439            if self.page_size_reduction
440            else None
441        )
442        reset_pagination = False
443        reduce_page_size = False
444        next_page_token = self._get_initial_next_page_token()
445        while True:
446            page_size_override = page_size_reducer.page_size_override if page_size_reducer else None
447            merged_records: MutableMapping[str, Any] = defaultdict(dict)
448            last_page_size = 0
449            last_record: Optional[Record] = None
450
451            response = None
452            try:
453                if self.additional_query_properties:
454                    for (
455                        properties
456                    ) in self.additional_query_properties.get_request_property_chunks():
457                        stream_slice = StreamSlice(
458                            partition=stream_slice.partition or {},
459                            cursor_slice=stream_slice.cursor_slice or {},
460                            extra_fields={"query_properties": properties},
461                        )
462                        response = self._fetch_next_page(
463                            stream_slice,
464                            next_page_token,
465                            **page_size_override_kwargs(page_size_override),
466                        )
467
468                        for current_record in records_generator_fn(response):
469                            if self.additional_query_properties.property_chunking:
470                                merge_key = self.additional_query_properties.property_chunking.get_merge_key(
471                                    current_record
472                                )
473                                if merge_key:
474                                    _deep_merge(merged_records[merge_key], current_record)
475                                else:
476                                    # We should still emit records even if the record did not have a merge key
477                                    pagination_tracker.observe(current_record)
478                                    last_page_size += 1
479                                    last_record = current_record
480                                    yield current_record
481                            else:
482                                pagination_tracker.observe(current_record)
483                                last_page_size += 1
484                                last_record = current_record
485                                yield current_record
486
487                    for merged_record in merged_records.values():
488                        record = Record(
489                            data=merged_record, stream_name=self.name, associated_slice=stream_slice
490                        )
491                        pagination_tracker.observe(record)
492                        last_page_size += 1
493                        last_record = record
494                        yield record
495                else:
496                    response = self._fetch_next_page(
497                        stream_slice,
498                        next_page_token,
499                        **page_size_override_kwargs(page_size_override),
500                    )
501                    for current_record in records_generator_fn(response):
502                        pagination_tracker.observe(current_record)
503                        last_page_size += 1
504                        last_record = current_record
505                        yield current_record
506            except PaginationResetRequiredException:
507                reset_pagination = True
508            except PageSizeReductionRequiredException as exception:
509                if page_size_reducer is None:
510                    # The action can be attached to a requester we cannot validate at config time, such as one
511                    # built by a custom error handler or a CustomRequester. Re-raise as the misconfiguration it
512                    # is: the exception being handled is the neutral "the API asked for a smaller page" signal.
513                    raise PageSizeReductionNotSupportedException(stream_name=self.name)
514                if last_page_size:
515                    # The reduction is safe only because it re-issues a page whose records were not emitted.
516                    # Every in-CDK way of reaching this raises from `_fetch_next_page`, before the record loop,
517                    # and the factory rejects the manifest constructs that would not, but a custom extractor,
518                    # filter or transformation can issue its own request from inside the record generator.
519                    # Re-issuing the page then duplicates the records already yielded, so fail instead.
520                    raise AirbyteTracedException(
521                        internal_message=f"Stream {self.name} requested a page size reduction after {last_page_size} records of the page had already been emitted",
522                        message=f"Stream {self.name} asked for a smaller page size in the middle of a page. The page cannot be requested again without duplicating the records already read from it. Move the REDUCE_PAGE_SIZE action to the error handler of the stream's main requester.",
523                        failure_type=FailureType.config_error,
524                    )
525                # Raises once the page size cannot be reduced any further and the retries allowed at that
526                # floor are spent, which is what stops the loop when the API keeps failing.
527                page_size_reducer.reduce(exception.error_message)
528                reduce_page_size = True
529            else:
530                if page_size_reducer:
531                    page_size_reducer.on_successful_page()
532                if not response:
533                    break
534
535            if reduce_page_size:
536                # Retry the very same page: neither the token nor the slice change, only the page size does -
537                # and not even that once the reducer is at its floor and only waiting is left.
538                reduce_page_size = False
539                continue
540
541            if reset_pagination or pagination_tracker.has_reached_limit():
542                next_page_token = self._get_initial_next_page_token()
543                previous_slice = stream_slice
544                stream_slice = pagination_tracker.reduce_slice_range_if_possible(
545                    stream_slice, original_stream_slice
546                )
547                LOGGER.info(
548                    f"Hitting PaginationReset event. StreamSlice used will go from {previous_slice} to {stream_slice}"
549                )
550                reset_pagination = False
551            else:
552                last_page_token_value = (
553                    next_page_token.get("next_page_token") if next_page_token else None
554                )
555                next_page_token = self._next_page_token(
556                    response=response,  # type:ignore # we are breaking from the loop on the try/else if there are no response so this should be fine
557                    last_page_size=last_page_size,
558                    last_record=last_record,
559                    last_page_token_value=last_page_token_value,
560                    **page_size_override_kwargs(page_size_override),
561                    **stream_slice_kwargs(self._next_page_token, stream_slice),
562                )
563                if response is not None and self.page_observer:
564                    self.page_observer(
565                        PageCompleted(
566                            stream_slice=stream_slice,
567                            response=response,
568                            pagination_exhausted=not next_page_token,
569                        )
570                    )
571                if not next_page_token:
572                    break
573
574        # Always return an empty generator just in case no records were ever yielded
575        yield from []
576
577    def _get_initial_next_page_token(self) -> Optional[Mapping[str, Any]]:
578        initial_token = self._paginator.get_initial_token()
579        next_page_token = {"next_page_token": initial_token} if initial_token is not None else None
580        return next_page_token
581
582    def read_records(
583        self,
584        records_schema: Mapping[str, Any],
585        stream_slice: Optional[StreamSlice] = None,
586    ) -> Iterable[StreamData]:
587        """
588        Fetch a stream's records from an HTTP API source
589
590        :param records_schema: json schema to describe record
591        :param stream_slice: The stream slice to read data for
592        :return: The records read from the API source
593        """
594        _slice = stream_slice or StreamSlice(partition={}, cursor_slice={})  # None-check
595        yield from self._read_records_or_split_request_window(
596            records_schema, _slice, _slice, depth=0
597        )
598
599    def _read_records_or_split_request_window(
600        self,
601        records_schema: Mapping[str, Any],
602        stream_slice: StreamSlice,
603        original_slice: StreamSlice,
604        depth: int,
605    ) -> Iterable[StreamData]:
606        """
607        Read `stream_slice` to completion, replacing it with smaller children and recursing into each of them in
608        turn when a `RequestWindowSplitRequiredException` is raised while reading it.
609
610        Each recursive call re-enters this method - and, through it, `_read_pages` - from scratch for the child
611        slice it is given, so every child gets its own paginator token, `PaginationTracker`, and
612        `PageSizeReducer` for free: nothing here needs to reset that state explicitly. Because the whole
613        recursion lives inside the single generator `DeclarativePartition.read()` consumes,
614        `PartitionReader.process_partition()` only calls `cursor.close_partition()` once, after this generator is
615        fully exhausted - so a failure anywhere in the recursion (a child that cannot be read, or a window that
616        cannot be split any further) propagates out without ever checkpointing the original partition.
617
618        `original_slice` is threaded through the recursion unchanged so every record yielded, however deep the
619        recursion went to produce it, is re-stamped with the partition's own slice rather than the child slice it
620        was actually read against - see `_reassociate_with_original_slice`. `depth` bounds the recursion
621        independently of `request_window_splitter`'s own no-progress guard: it is enforced here so a
622        misbehaving custom cursor cannot recurse indefinitely regardless of what that guard does or does not
623        catch.
624        """
625        record_generator = partial(
626            self._parse_records,
627            stream_slice=stream_slice,
628            records_schema=records_schema,
629        )
630
631        emitted_count = 0
632        try:
633            records: Iterable[Mapping[str, Any]] = self._read_pages(record_generator, stream_slice)
634            if self.post_pagination_filter:
635                # A data feed paginates until it reaches a record older than the cursor, so the page that triggers the stop
636                # condition still holds already-synced records. Those are filtered here rather than in the record selector
637                # so that the paginator keeps seeing the whole page: the stop condition is evaluated on the last record of
638                # the page, which is precisely one of the records being dropped. Two consequences of filtering this late:
639                # the pagination tracker observes the dropped records, and a `file_uploader` on the record selector has
640                # already uploaded their files by the time they are dropped.
641                records = self.post_pagination_filter.filter_records(
642                    records,
643                    # the filter is only used for its cursor comparison, which does not read the stream state
644                    stream_state={},
645                    stream_slice=stream_slice,
646                )
647            for record in records:
648                emitted_count += 1
649                yield self._reassociate_with_original_slice(record, original_slice)
650        except RequestWindowSplitRequiredException as exception:
651            if self.request_window_splitting is None or self.request_window_splitter is None:
652                raise RequestWindowSplitNotSupportedException(stream_name=self.name) from exception
653
654            if emitted_count:
655                # Splitting and re-reading the window re-emits these records: a failed partition is never
656                # checkpointed, so the next attempt would re-emit them anyway. Always replaying makes forward
657                # progress instead of retrying the same oversized window forever; primary-key dedup downstream
658                # handles the duplicates either way.
659                LOGGER.warning(
660                    f"Stream {self.name} already emitted {emitted_count} record(s) from {stream_slice} before "
661                    f"it was rejected; splitting and re-reading the window may re-emit them."
662                )
663
664            rejection = f": {exception.error_message}" if exception.error_message else ""
665            if depth >= _MAX_REQUEST_WINDOW_SPLIT_DEPTH:
666                # Logged separately from the exception below so it stays visible even if the trace message's
667                # internal_message isn't surfaced by whatever catches it.
668                LOGGER.warning(
669                    f"Stream {self.name} hit the maximum request window split depth "
670                    f"({_MAX_REQUEST_WINDOW_SPLIT_DEPTH}) while splitting {original_slice}."
671                )
672                raise AirbyteTracedException(
673                    internal_message=f"Stream {self.name} exceeded the maximum request window split depth of {_MAX_REQUEST_WINDOW_SPLIT_DEPTH} while splitting {original_slice}; the API still rejected {stream_slice}{rejection}",
674                    # `transient_error`, so the only remediation is the connector's own, if it defined one.
675                    message=self._with_request_window_failure_message(
676                        f"Stream {self.name} could not split its request window to a size the API accepts "
677                        f"within {_MAX_REQUEST_WINDOW_SPLIT_DEPTH} splits."
678                    ),
679                    failure_type=FailureType.transient_error,
680                ) from exception
681
682            children = self.request_window_splitter(
683                stream_slice, self.request_window_splitting.min_split_window
684            )
685            if children is None:
686                min_split_window_note = (
687                    f", its configured `min_split_window` ({self.request_window_splitting.min_split_window})"
688                    if self.request_window_splitting.min_split_window
689                    else ""
690                )
691                raise AirbyteTracedException(
692                    internal_message=f"Stream {self.name} could not split its request window {stream_slice} any further; the API rejected it{rejection}",
693                    # `transient_error` regardless of how the triggering response was classified: exhaustion is
694                    # never the user's fault, matching `PageSizeReducer`'s equivalent exhaustion branch.
695                    message=self._with_request_window_failure_message(
696                        f"The API kept rejecting stream {self.name}'s request window even at the smallest window "
697                        f"its cursor allows{min_split_window_note}, or splitting the window further would not "
698                        f"make progress."
699                    ),
700                    failure_type=FailureType.transient_error,
701                ) from exception
702
703            LOGGER.info(
704                f"Stream {self.name}: the API rejected request window {stream_slice} (split depth {depth}); "
705                f"reducing it to {children} and reading each in turn."
706            )
707            for child in children:
708                yield from self._read_records_or_split_request_window(
709                    records_schema, child, original_slice, depth + 1
710                )
711
712    def _with_request_window_failure_message(self, message: str) -> str:
713        """
714        :return: the message followed by `request_window_splitting`'s configured `failure_message`, when it
715            defined one - mirroring `PageSizeReducer._with_failure_message`
716        """
717        assert self.request_window_splitting is not None
718        failure_message = (self.request_window_splitting.failure_message or "").strip()
719        if not failure_message:
720            return message
721        return f"{message} {failure_message}"
722
723    @staticmethod
724    def _reassociate_with_original_slice(
725        record: StreamData, original_slice: StreamSlice
726    ) -> StreamData:
727        """
728        Records read while recursing into a split child window are parsed against that child's `StreamSlice`,
729        so `RecordSelector` stamps them with the child as `associated_slice`. `DeclarativePartition.read()`
730        passes an already-built `Record` through unchanged rather than re-wrapping it, so left uncorrected the
731        child slice - not the partition's own slice - is what `ConcurrentCursor.observe()` would key its
732        per-partition bookkeeping by. `close_partition()` always looks that bookkeeping up by the partition's own
733        slice, so it would never find it, silently losing the precise most-recently-observed cursor value for a
734        split partition (recoverable in practice today only because callers of that lookup fall back to the
735        slice's end boundary when it is missing). Re-stamping every record with `original_slice` here restores
736        the same tracking a non-split read already gets, rather than relying on that fallback.
737        """
738        if isinstance(record, Record) and record.associated_slice is not original_slice:
739            return Record(
740                data=record.data,
741                stream_name=record.stream_name,
742                associated_slice=original_slice,
743                file_reference=record.file_reference,
744            )
745        return record
746
747    def _parse_records(
748        self,
749        response: Optional[requests.Response],
750        records_schema: Mapping[str, Any],
751        stream_slice: Optional[StreamSlice],
752    ) -> Iterable[Record]:
753        yield from self._parse_response(
754            response,
755            stream_slice=stream_slice,
756            records_schema=records_schema,
757        )
758
759    def must_deduplicate_query_params(self) -> bool:
760        return True
761
762    @staticmethod
763    def _to_partition_key(to_serialize: Any) -> str:
764        # separators have changed in Python 3.4. To avoid being impacted by further change, we explicitly specify our own value
765        return json.dumps(to_serialize, indent=None, separators=(",", ":"), sort_keys=True)

 19@dataclass
 20class AsyncRetriever(Retriever):
 21    config: Config
 22    parameters: InitVar[Mapping[str, Any]]
 23    record_selector: RecordSelector
 24    stream_slicer: AsyncJobPartitionRouter
 25    slice_logger: AlwaysLogSliceLogger = field(
 26        init=False,
 27        default_factory=lambda: AlwaysLogSliceLogger(),
 28    )
 29
 30    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
 31        self._parameters = parameters
 32
 33    @property
 34    def exit_on_rate_limit(self) -> bool:
 35        """
 36        Whether to exit on rate limit. This is a property of the job repository
 37        and not the stream slicer. The stream slicer is responsible for creating
 38        the jobs, but the job repository is responsible for managing the rate
 39        limits and other job-related properties.
 40
 41        Note:
 42         - If the `creation_requester` cannot place / create the job - it might be the case of the RateLimits
 43         - If the `creation_requester` can place / create the job - it means all other requesters should successfully manage
 44           to complete the results.
 45        """
 46        job_orchestrator = self.stream_slicer._job_orchestrator
 47        if job_orchestrator is None:
 48            # Default value when orchestrator is not available
 49            return False
 50        return job_orchestrator._job_repository.creation_requester.exit_on_rate_limit  # type: ignore
 51
 52    @exit_on_rate_limit.setter
 53    def exit_on_rate_limit(self, value: bool) -> None:
 54        """
 55        Sets the `exit_on_rate_limit` property of the job repository > creation_requester,
 56        meaning that the Job cannot be placed / created if the rate limit is reached.
 57        Thus no further work on managing jobs is expected to be done.
 58        """
 59        job_orchestrator = self.stream_slicer._job_orchestrator
 60        if job_orchestrator is not None:
 61            job_orchestrator._job_repository.creation_requester.exit_on_rate_limit = value  # type: ignore[attr-defined, assignment]
 62
 63    def _validate_and_get_stream_slice_jobs(
 64        self, stream_slice: Optional[StreamSlice] = None
 65    ) -> Iterable[AsyncJob]:
 66        """
 67        Validates the stream_slice argument and returns the partition from it.
 68
 69        Args:
 70            stream_slice (Optional[StreamSlice]): The stream slice to validate and extract the partition from.
 71
 72        Returns:
 73            AsyncPartition: The partition extracted from the stream_slice.
 74
 75        Raises:
 76            AirbyteTracedException: If the stream_slice is not an instance of StreamSlice or if the partition is not present in the stream_slice.
 77
 78        """
 79        return stream_slice.extra_fields.get("jobs", []) if stream_slice else []
 80
 81    def read_records(
 82        self,
 83        records_schema: Mapping[str, Any],
 84        stream_slice: Optional[StreamSlice] = None,
 85    ) -> Iterable[StreamData]:
 86        if stream_slice is None:
 87            # Slice-less callers such as DynamicSchemaLoader read this way: create and poll the jobs here
 88            for job_slice in self.stream_slicer.stream_slices():
 89                yield from self.read_records(records_schema, job_slice)
 90            return
 91
 92        # emit the slice_descriptor log message, for connector builder TestRead
 93        yield self.slice_logger.create_slice_log_message(stream_slice.cursor_slice)  # type: ignore
 94
 95        jobs: Iterable[AsyncJob] = self._validate_and_get_stream_slice_jobs(stream_slice)
 96        records: Iterable[Mapping[str, Any]] = self.stream_slicer.fetch_records(jobs)
 97
 98        yield from self.record_selector.filter_and_transform(
 99            all_data=records,
100            stream_state={},  # stream_state as an interpolation context is deprecated
101            records_schema=records_schema,
102            stream_slice=stream_slice,
103        )

788@deprecated(
789    "This class is experimental. Use at your own risk.",
790    category=ExperimentalClassWarning,
791)
792@dataclass
793class LazySimpleRetriever(SimpleRetriever):
794    """
795    A retriever that supports lazy loading from parent streams.
796    """
797
798    def _read_pages(
799        self,
800        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
801        stream_slice: StreamSlice,
802    ) -> Iterable[Record]:
803        response = stream_slice.extra_fields["child_response"]
804        if response:
805            last_page_size, last_record = 0, None
806            for record in records_generator_fn(response):  # type: ignore[call-arg] # only _parse_records expected as a func
807                last_page_size += 1
808                last_record = record
809                yield record
810
811            next_page_token = self._next_page_token(
812                response,
813                last_page_size,
814                last_record,
815                None,
816                **stream_slice_kwargs(self._next_page_token, stream_slice),
817            )
818            if next_page_token:
819                yield from self._paginate(
820                    next_page_token,
821                    records_generator_fn,
822                    stream_slice,
823                )
824
825            yield from []
826        else:
827            # coderabbit detected an interesting bug/gap where if we were to not get a child_response, we
828            # might recurse forever. This might not be the case, but it is worth noting that this code path
829            # isn't comprehensively tested.
830            yield from self._read_pages(records_generator_fn, stream_slice)
831
832    def _paginate(
833        self,
834        next_page_token: Any,
835        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
836        stream_slice: StreamSlice,
837    ) -> Iterable[Record]:
838        """Handle pagination by fetching subsequent pages."""
839        pagination_complete = False
840
841        while not pagination_complete:
842            response = self._fetch_next_page(stream_slice, next_page_token)
843            last_page_size, last_record = 0, None
844
845            for record in records_generator_fn(response):  # type: ignore[call-arg] # only _parse_records expected as a func
846                last_page_size += 1
847                last_record = record
848                yield record
849
850            if not response:
851                pagination_complete = True
852            else:
853                last_page_token_value = (
854                    next_page_token.get("next_page_token") if next_page_token else None
855                )
856                next_page_token = self._next_page_token(
857                    response,
858                    last_page_size,
859                    last_record,
860                    last_page_token_value,
861                    **stream_slice_kwargs(self._next_page_token, stream_slice),
862                )
863
864                if not next_page_token:
865                    pagination_complete = True