airbyte_cdk.sources.declarative.retrievers

15class Retriever:
16    """
17    Responsible for fetching a stream's records from an HTTP API source.
18    """
19
20    @abstractmethod
21    def read_records(
22        self,
23        records_schema: Mapping[str, Any],
24        stream_slice: Optional[StreamSlice] = None,
25    ) -> Iterable[StreamData]:
26        """
27        Fetch a stream's records from an HTTP API source
28
29        :param records_schema: json schema to describe record
30        :param stream_slice: The stream slice to read data for
31        :return: The records read from the API source
32        """
33
34    @deprecated("Stream slicing is being moved to the stream level.")
35    def stream_slices(self) -> Iterable[Optional[StreamSlice]]:
36        """Does nothing as this method is deprecated, so underlying Retriever implementations
37        do not need to implement this.
38        """
39        yield from []
40
41    @property
42    @deprecated("State management is being moved to the stream level.")
43    def state(self) -> StreamState:
44        """
45        Does nothing as this method is deprecated, so underlying Retriever implementations
46        do not need to implement this.
47        """
48        return {}
49
50    @state.setter
51    @deprecated("State management is being moved to the stream level.")
52    def state(self, value: StreamState) -> None:
53        """
54        Does nothing as this method is deprecated, so underlying Retriever implementations
55        do not need to implement this.
56        """
57        pass

 58@dataclass
 59class SimpleRetriever(Retriever):
 60    """
 61    Retrieves records by synchronously sending requests to fetch records.
 62
 63    The retriever acts as an orchestrator between the requester, the record selector, the paginator, and the stream slicer.
 64
 65    For each stream slice, submit requests until there are no more pages of records to fetch.
 66
 67    This retriever currently inherits from HttpStream to reuse the request submission and pagination machinery.
 68    As a result, some of the parameters passed to some methods are unused.
 69    The two will be decoupled in a future release.
 70
 71    Attributes:
 72        stream_name (str): The stream's name
 73        stream_primary_key (Optional[Union[str, List[str], List[List[str]]]]): The stream's primary key
 74        requester (Requester): The HTTP requester
 75        record_selector (HttpSelector): The record selector
 76        paginator (Optional[Paginator]): The paginator
 77        stream_slicer (Optional[StreamSlicer]): The stream slicer
 78        parameters (Mapping[str, Any]): Additional runtime parameters to be used for string interpolation
 79        post_pagination_filter (Optional[ClientSideIncrementalRecordFilterDecorator]): Set for data feed streams only.
 80            Records the cursor considers already synced are dropped once pagination has observed them
 81    """
 82
 83    requester: Requester
 84    record_selector: HttpSelector
 85    config: Config
 86    parameters: InitVar[Mapping[str, Any]]
 87    name: str
 88    _name: Union[InterpolatedString, str] = field(init=False, repr=False, default="")
 89    primary_key: Optional[Union[str, List[str], List[List[str]]]]
 90    _primary_key: str = field(init=False, repr=False, default="")
 91    paginator: Optional[Paginator] = None
 92    stream_slicer: StreamSlicer = field(
 93        default_factory=lambda: SinglePartitionRouter(parameters={})
 94    )
 95    request_option_provider: RequestOptionsProvider = field(
 96        default_factory=lambda: DefaultRequestOptionsProvider(parameters={})
 97    )
 98    ignore_stream_slicer_parameters_on_paginated_requests: bool = False
 99    additional_query_properties: Optional[QueryProperties] = None
100    log_formatter: Optional[Callable[[requests.Response], Any]] = None
101    pagination_tracker_factory: Callable[[], PaginationTracker] = field(
102        default_factory=lambda: lambda: PaginationTracker()
103    )
104    post_pagination_filter: Optional[ClientSideIncrementalRecordFilterDecorator] = None
105
106    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
107        self._paginator = self.paginator or NoPagination(parameters=parameters)
108        self._parameters = parameters
109        self._name = (
110            InterpolatedString(self._name, parameters=parameters)
111            if isinstance(self._name, str)
112            else self._name
113        )
114
115    @property  # type: ignore
116    def name(self) -> str:
117        """
118        :return: Stream name
119        """
120        return (
121            str(self._name.eval(self.config))
122            if isinstance(self._name, InterpolatedString)
123            else self._name
124        )
125
126    @name.setter
127    def name(self, value: str) -> None:
128        if not isinstance(value, property):
129            self._name = value
130
131    def _get_mapping(
132        self, method: Callable[..., Optional[Union[Mapping[str, Any], str]]], **kwargs: Any
133    ) -> Tuple[Union[Mapping[str, Any], str], Set[str]]:
134        """
135        Get mapping from the provided method, and get the keys of the mapping.
136        If the method returns a string, it will return the string and an empty set.
137        If the method returns a dict, it will return the dict and its keys.
138        """
139        mapping = method(**kwargs) or {}
140        keys = set(mapping.keys()) if not isinstance(mapping, str) else set()
141        return mapping, keys
142
143    def _get_request_options(
144        self,
145        stream_slice: Optional[StreamSlice],
146        next_page_token: Optional[Mapping[str, Any]],
147        paginator_method: Callable[..., Optional[Union[Mapping[str, Any], str]]],
148        stream_slicer_method: Callable[..., Optional[Union[Mapping[str, Any], str]]],
149    ) -> Union[Mapping[str, Any], str]:
150        """
151        Get the request_option from the paginator and the stream slicer.
152        Raise a ValueError if there's a key collision
153        Returned merged mapping otherwise
154        """
155        is_body_json = paginator_method.__name__ == "get_request_body_json"
156
157        mappings = [
158            paginator_method(
159                stream_slice=stream_slice,
160                next_page_token=next_page_token,
161            ),
162        ]
163        if not next_page_token or not self.ignore_stream_slicer_parameters_on_paginated_requests:
164            mappings.append(
165                stream_slicer_method(
166                    stream_slice=stream_slice,
167                    next_page_token=next_page_token,
168                )
169            )
170        return combine_mappings(mappings, allow_same_value_merge=is_body_json)
171
172    def _request_headers(
173        self,
174        stream_slice: Optional[StreamSlice] = None,
175        next_page_token: Optional[Mapping[str, Any]] = None,
176    ) -> Mapping[str, Any]:
177        """
178        Specifies request headers.
179        Authentication headers will overwrite any overlapping headers returned from this method.
180        """
181        headers = self._get_request_options(
182            stream_slice,
183            next_page_token,
184            self._paginator.get_request_headers,
185            self.request_option_provider.get_request_headers,
186        )
187        if isinstance(headers, str):
188            raise ValueError("Request headers cannot be a string")
189        return {str(k): str(v) for k, v in headers.items()}
190
191    def _request_params(
192        self,
193        stream_slice: Optional[StreamSlice] = None,
194        next_page_token: Optional[Mapping[str, Any]] = None,
195    ) -> Mapping[str, Any]:
196        """
197        Specifies the query parameters that should be set on an outgoing HTTP request given the inputs.
198
199        E.g: you might want to define query parameters for paging if next_page_token is not None.
200        """
201        params = self._get_request_options(
202            stream_slice,
203            next_page_token,
204            self._paginator.get_request_params,
205            self.request_option_provider.get_request_params,
206        )
207        if isinstance(params, str):
208            raise ValueError("Request params cannot be a string")
209        return params
210
211    def _request_body_data(
212        self,
213        stream_slice: Optional[StreamSlice] = None,
214        next_page_token: Optional[Mapping[str, Any]] = None,
215    ) -> Union[Mapping[str, Any], str]:
216        """
217        Specifies how to populate the body of the request with a non-JSON payload.
218
219        If returns a ready text that it will be sent as is.
220        If returns a dict that it will be converted to a urlencoded form.
221        E.g. {"key1": "value1", "key2": "value2"} => "key1=value1&key2=value2"
222
223        At the same time only one of the 'request_body_data' and 'request_body_json' functions can be overridden.
224        """
225        return self._get_request_options(
226            stream_slice,
227            next_page_token,
228            self._paginator.get_request_body_data,
229            self.request_option_provider.get_request_body_data,
230        )
231
232    def _request_body_json(
233        self,
234        stream_slice: Optional[StreamSlice] = None,
235        next_page_token: Optional[Mapping[str, Any]] = None,
236    ) -> Optional[Mapping[str, Any]]:
237        """
238        Specifies how to populate the body of the request with a JSON payload.
239
240        At the same time only one of the 'request_body_data' and 'request_body_json' functions can be overridden.
241        """
242        body_json = self._get_request_options(
243            stream_slice,
244            next_page_token,
245            self._paginator.get_request_body_json,
246            self.request_option_provider.get_request_body_json,
247        )
248        if isinstance(body_json, str):
249            raise ValueError("Request body json cannot be a string")
250        return body_json
251
252    def _paginator_path(
253        self,
254        next_page_token: Optional[Mapping[str, Any]] = None,
255        stream_slice: Optional[StreamSlice] = None,
256    ) -> Optional[str]:
257        """
258        If the paginator points to a path, follow it, else return nothing so the requester is used.
259        :param next_page_token:
260        :return:
261        """
262        return self._paginator.path(
263            next_page_token=next_page_token,
264            stream_state={},  # stream_state as an interpolation context is deprecated
265            stream_slice=stream_slice,
266        )
267
268    def _parse_response(
269        self,
270        response: Optional[requests.Response],
271        records_schema: Mapping[str, Any],
272        stream_slice: Optional[StreamSlice] = None,
273        next_page_token: Optional[Mapping[str, Any]] = None,
274    ) -> Iterable[Record]:
275        if not response:
276            yield from []
277        else:
278            yield from self.record_selector.select_records(
279                response=response,
280                stream_state={},  # stream_state as an interpolation context is deprecated
281                records_schema=records_schema,
282                stream_slice=stream_slice,
283                next_page_token=next_page_token,
284            )
285
286    @property  # type: ignore
287    def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
288        """The stream's primary key"""
289        return self._primary_key
290
291    @primary_key.setter
292    def primary_key(self, value: str) -> None:
293        if not isinstance(value, property):
294            self._primary_key = value
295
296    def _next_page_token(
297        self,
298        response: requests.Response,
299        last_page_size: int,
300        last_record: Optional[Record],
301        last_page_token_value: Optional[Any],
302    ) -> Optional[Mapping[str, Any]]:
303        """
304        Specifies a pagination strategy.
305
306        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.
307
308        :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.
309        """
310        return self._paginator.next_page_token(
311            response=response,
312            last_page_size=last_page_size,
313            last_record=last_record,
314            last_page_token_value=last_page_token_value,
315        )
316
317    def _fetch_next_page(
318        self,
319        stream_slice: StreamSlice,
320        next_page_token: Optional[Mapping[str, Any]] = None,
321    ) -> Optional[requests.Response]:
322        return self.requester.send_request(
323            path=self._paginator_path(
324                next_page_token=next_page_token,
325                stream_slice=stream_slice,
326            ),
327            stream_state={},  # stream_state as an interpolation context is deprecated
328            stream_slice=stream_slice,
329            next_page_token=next_page_token,
330            request_headers=self._request_headers(
331                stream_slice=stream_slice,
332                next_page_token=next_page_token,
333            ),
334            request_params=self._request_params(
335                stream_slice=stream_slice,
336                next_page_token=next_page_token,
337            ),
338            request_body_data=self._request_body_data(
339                stream_slice=stream_slice,
340                next_page_token=next_page_token,
341            ),
342            request_body_json=self._request_body_json(
343                stream_slice=stream_slice,
344                next_page_token=next_page_token,
345            ),
346            log_formatter=self.log_formatter,
347        )
348
349    # This logic is similar to _read_pages in the HttpStream class. When making changes here, consider making changes there as well.
350    def _read_pages(
351        self,
352        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
353        stream_slice: StreamSlice,
354    ) -> Iterable[Record]:
355        original_stream_slice = stream_slice
356        pagination_tracker = self.pagination_tracker_factory()
357        reset_pagination = False
358        next_page_token = self._get_initial_next_page_token()
359        while True:
360            merged_records: MutableMapping[str, Any] = defaultdict(dict)
361            last_page_size = 0
362            last_record: Optional[Record] = None
363
364            response = None
365            try:
366                if self.additional_query_properties:
367                    for (
368                        properties
369                    ) in self.additional_query_properties.get_request_property_chunks():
370                        stream_slice = StreamSlice(
371                            partition=stream_slice.partition or {},
372                            cursor_slice=stream_slice.cursor_slice or {},
373                            extra_fields={"query_properties": properties},
374                        )
375                        response = self._fetch_next_page(stream_slice, next_page_token)
376
377                        for current_record in records_generator_fn(response):
378                            if self.additional_query_properties.property_chunking:
379                                merge_key = self.additional_query_properties.property_chunking.get_merge_key(
380                                    current_record
381                                )
382                                if merge_key:
383                                    _deep_merge(merged_records[merge_key], current_record)
384                                else:
385                                    # We should still emit records even if the record did not have a merge key
386                                    pagination_tracker.observe(current_record)
387                                    last_page_size += 1
388                                    last_record = current_record
389                                    yield current_record
390                            else:
391                                pagination_tracker.observe(current_record)
392                                last_page_size += 1
393                                last_record = current_record
394                                yield current_record
395
396                    for merged_record in merged_records.values():
397                        record = Record(
398                            data=merged_record, stream_name=self.name, associated_slice=stream_slice
399                        )
400                        pagination_tracker.observe(record)
401                        last_page_size += 1
402                        last_record = record
403                        yield record
404                else:
405                    response = self._fetch_next_page(stream_slice, next_page_token)
406                    for current_record in records_generator_fn(response):
407                        pagination_tracker.observe(current_record)
408                        last_page_size += 1
409                        last_record = current_record
410                        yield current_record
411            except PaginationResetRequiredException:
412                reset_pagination = True
413            else:
414                if not response:
415                    break
416
417            if reset_pagination or pagination_tracker.has_reached_limit():
418                next_page_token = self._get_initial_next_page_token()
419                previous_slice = stream_slice
420                stream_slice = pagination_tracker.reduce_slice_range_if_possible(
421                    stream_slice, original_stream_slice
422                )
423                LOGGER.info(
424                    f"Hitting PaginationReset event. StreamSlice used will go from {previous_slice} to {stream_slice}"
425                )
426                reset_pagination = False
427            else:
428                last_page_token_value = (
429                    next_page_token.get("next_page_token") if next_page_token else None
430                )
431                next_page_token = self._next_page_token(
432                    response=response,  # type:ignore # we are breaking from the loop on the try/else if there are no response so this should be fine
433                    last_page_size=last_page_size,
434                    last_record=last_record,
435                    last_page_token_value=last_page_token_value,
436                )
437                if not next_page_token:
438                    break
439
440        # Always return an empty generator just in case no records were ever yielded
441        yield from []
442
443    def _get_initial_next_page_token(self) -> Optional[Mapping[str, Any]]:
444        initial_token = self._paginator.get_initial_token()
445        next_page_token = {"next_page_token": initial_token} if initial_token is not None else None
446        return next_page_token
447
448    def read_records(
449        self,
450        records_schema: Mapping[str, Any],
451        stream_slice: Optional[StreamSlice] = None,
452    ) -> Iterable[StreamData]:
453        """
454        Fetch a stream's records from an HTTP API source
455
456        :param records_schema: json schema to describe record
457        :param stream_slice: The stream slice to read data for
458        :return: The records read from the API source
459        """
460        _slice = stream_slice or StreamSlice(partition={}, cursor_slice={})  # None-check
461
462        record_generator = partial(
463            self._parse_records,
464            stream_slice=stream_slice,
465            records_schema=records_schema,
466        )
467        records: Iterable[Mapping[str, Any]] = self._read_pages(record_generator, _slice)
468        if self.post_pagination_filter:
469            # A data feed paginates until it reaches a record older than the cursor, so the page that triggers the stop
470            # condition still holds already-synced records. Those are filtered here rather than in the record selector
471            # so that the paginator keeps seeing the whole page: the stop condition is evaluated on the last record of
472            # the page, which is precisely one of the records being dropped. Two consequences of filtering this late:
473            # the pagination tracker observes the dropped records, and a `file_uploader` on the record selector has
474            # already uploaded their files by the time they are dropped.
475            records = self.post_pagination_filter.filter_records(
476                records,
477                # the filter is only used for its cursor comparison, which does not read the stream state
478                stream_state={},
479                stream_slice=_slice,
480            )
481        yield from records
482
483    def _parse_records(
484        self,
485        response: Optional[requests.Response],
486        records_schema: Mapping[str, Any],
487        stream_slice: Optional[StreamSlice],
488    ) -> Iterable[Record]:
489        yield from self._parse_response(
490            response,
491            stream_slice=stream_slice,
492            records_schema=records_schema,
493        )
494
495    def must_deduplicate_query_params(self) -> bool:
496        return True
497
498    @staticmethod
499    def _to_partition_key(to_serialize: Any) -> str:
500        # separators have changed in Python 3.4. To avoid being impacted by further change, we explicitly specify our own value
501        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        # emit the slice_descriptor log message, for connector builder TestRead
87        yield self.slice_logger.create_slice_log_message(stream_slice.cursor_slice)  # type: ignore
88
89        jobs: Iterable[AsyncJob] = self._validate_and_get_stream_slice_jobs(stream_slice)
90        records: Iterable[Mapping[str, Any]] = self.stream_slicer.fetch_records(jobs)
91
92        yield from self.record_selector.filter_and_transform(
93            all_data=records,
94            stream_state={},  # stream_state as an interpolation context is deprecated
95            records_schema=records_schema,
96            stream_slice=stream_slice,
97        )

524@deprecated(
525    "This class is experimental. Use at your own risk.",
526    category=ExperimentalClassWarning,
527)
528@dataclass
529class LazySimpleRetriever(SimpleRetriever):
530    """
531    A retriever that supports lazy loading from parent streams.
532    """
533
534    def _read_pages(
535        self,
536        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
537        stream_slice: StreamSlice,
538    ) -> Iterable[Record]:
539        response = stream_slice.extra_fields["child_response"]
540        if response:
541            last_page_size, last_record = 0, None
542            for record in records_generator_fn(response):  # type: ignore[call-arg] # only _parse_records expected as a func
543                last_page_size += 1
544                last_record = record
545                yield record
546
547            next_page_token = self._next_page_token(response, last_page_size, last_record, None)
548            if next_page_token:
549                yield from self._paginate(
550                    next_page_token,
551                    records_generator_fn,
552                    stream_slice,
553                )
554
555            yield from []
556        else:
557            # coderabbit detected an interesting bug/gap where if we were to not get a child_response, we
558            # might recurse forever. This might not be the case, but it is worth noting that this code path
559            # isn't comprehensively tested.
560            yield from self._read_pages(records_generator_fn, stream_slice)
561
562    def _paginate(
563        self,
564        next_page_token: Any,
565        records_generator_fn: Callable[[Optional[requests.Response]], Iterable[Record]],
566        stream_slice: StreamSlice,
567    ) -> Iterable[Record]:
568        """Handle pagination by fetching subsequent pages."""
569        pagination_complete = False
570
571        while not pagination_complete:
572            response = self._fetch_next_page(stream_slice, next_page_token)
573            last_page_size, last_record = 0, None
574
575            for record in records_generator_fn(response):  # type: ignore[call-arg] # only _parse_records expected as a func
576                last_page_size += 1
577                last_record = record
578                yield record
579
580            if not response:
581                pagination_complete = True
582            else:
583                last_page_token_value = (
584                    next_page_token.get("next_page_token") if next_page_token else None
585                )
586                next_page_token = self._next_page_token(
587                    response, last_page_size, last_record, last_page_token_value
588                )
589
590                if not next_page_token:
591                    pagination_complete = True