fix(low-code): stop OffsetIncrement/PageIncrement re-reading pages

airbyte_cdk/sources/declarative/extractors/combined_extractor.py

@@ -34,9 +34,8 @@ class _NoRecord:
 class _DroppedEmptyRecord:
     """Stands in for an empty record dropped by `skip_empty_records` when the extractor counts a page.
 
-    A record-counting paginator only takes the length of what the extractor yields, so the
-    placeholder makes the dropped record count toward the page size without ever reaching the
-    record stream.
+    `RecordSelector` counts every item the extractor yields and discards this placeholder, so the
+    dropped record counts toward the page size without ever reaching the record stream.
     """
 
 
@@ -132,12 +131,10 @@ class CombinedExtractor(RecordExtractor):
     Combining three extractors over a large response costs roughly three times the decode time of
     a single one.
 
-    Note that an `OffsetIncrement` or `PageIncrement` paginator builds its own copy of the
-    extractor to count the records of a page, which doubles that cost.
-
     ## Record-counting paginators
 
-    The count those paginators obtain is the combined count, not the API's page size:
+    The count `OffsetIncrement` and `PageIncrement` obtain is the combined count, not the API's
+    page size:
 
     - `union` returns the SUM over all sub-extractors. An `OffsetIncrement` advances the offset by
       that sum, so the next page starts past the records that were never read — two sub-extractors
@@ -154,19 +151,19 @@ class CombinedExtractor(RecordExtractor):
     Empty records dropped by `skip_empty_records` are part of that count. The API returned them as
     slots of the page, so a full page holding a null must still read as full: counting only the
     surviving records would make both paginators stop at the first page with a null, and would
-    make an `OffsetIncrement` re-request records it already read. The factory builds the
-    paginator's copy of the extractor with `count_dropped_empty_records` set, which makes that copy
-    yield a placeholder for every empty record it drops. The winner under `first_match` is still
-    chosen by the surviving records alone, so the paginator counts the same sub-extractor the
-    record stream reads.
+    make an `OffsetIncrement` re-request records it already read. The factory sets
+    `count_dropped_empty_records` on the extractor of a `RecordSelector`, which makes it yield a
+    placeholder for every empty record it drops; the `RecordSelector` counts the placeholders and
+    discards them. The winner under `first_match` is still chosen by the surviving records alone,
+    so the placeholders never change which records are emitted.
 
     Attributes:
         extractors (List[RecordExtractor]): The sub-extractors to combine. At least one is required.
         mode (CombineMode): How the sub-extractor outputs are combined. Defaults to `union`.
         skip_empty_records (bool): Whether falsy records are dropped before they are combined.
             Defaults to `False`.
-        count_dropped_empty_records (bool): Set by the factory on the copy a record-counting
-            paginator uses, never on the one that emits records: yield a placeholder in place of
+        count_dropped_empty_records (bool): Set by the factory on the extractor of a
+            `RecordSelector`, which discards the placeholders: yield a placeholder in place of
             every dropped empty record so it counts toward the page size. Defaults to `False`.
     """
 
@@ -228,14 +225,11 @@ def _without_empty_records(
         dropped = 0
         for record in records:
             if not record:
+                dropped += 1
                 if self.count_dropped_empty_records:
                     yield _DROPPED_EMPTY_RECORD  # type: ignore[misc]  # only ever counted, never emitted
-                else:
-                    dropped += 1
                 continue
             yield record
-        # The paginator's copy re-reads the page the record stream already read, so only the
-        # record stream reports the drop.
         if dropped:
             logger.warning(
                 "CombinedExtractor dropped %s empty record(s) yielded by sub-extractor %s (%s) "

airbyte_cdk/sources/declarative/extractors/record_selector.py

@@ -7,6 +7,7 @@
 
 import requests
 
+from airbyte_cdk.sources.declarative.extractors.combined_extractor import _DroppedEmptyRecord
 from airbyte_cdk.sources.declarative.extractors.http_selector import HttpSelector
 from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor
 from airbyte_cdk.sources.declarative.extractors.record_filter import RecordFilter
@@ -20,6 +21,19 @@
 from airbyte_cdk.sources.types import Config, Record, StreamSlice, StreamState
 from airbyte_cdk.sources.utils.transform import TypeTransformer
 
+_EXTRACTED_RECORD_COUNT = "_airbyte_extracted_record_count"
+
+
+def extracted_record_count(response: requests.Response, default: int) -> int:
+    """How many records the extractor of a `RecordSelector` returned for `response` before any of
+    them was filtered out, or `default` if no `RecordSelector` has read the whole response.
+
+    Record-counting paginators compare this count with the page size: once a record filter dropped
+    records from a full page, the number of records emitted looks like a short last page.
+    """
+    count = getattr(response, _EXTRACTED_RECORD_COUNT, None)
+    return count if isinstance(count, int) else default
+
 
 @dataclass
 class RecordSelector(HttpSelector):
@@ -86,10 +100,24 @@ def select_records(
         :param next_page_token: The paginator token
         :return: List of Records selected from the response
         """
-        all_data: Iterable[Mapping[str, Any]] = self.extractor.extract_records(response)
+        all_data: Iterable[Mapping[str, Any]] = self._extract_and_count(response)
         yield from self.filter_and_transform(
             all_data, stream_state, records_schema, stream_slice, next_page_token
         )
+        # A custom record filter can stop reading before the end of the page, which would leave the
+        # page uncounted.
+        for _ in all_data:
+            pass
+
+    def _extract_and_count(self, response: requests.Response) -> Iterable[Mapping[str, Any]]:
+        count = 0
+        for record in self.extractor.extract_records(response):
+            count += 1
+            if not isinstance(record, _DroppedEmptyRecord):
+                yield record
+        # Kept on the response, not on the selector: a stream has one retriever, and so one
+        # selector, reading all of its partitions concurrently.
+        setattr(response, _EXTRACTED_RECORD_COUNT, count)
 
     def filter_and_transform(
         self,

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -3316,32 +3316,13 @@ def _combined_extractor_tree_contains_union(model: CombinedExtractorModel) -> bo
             for sub_extractor in model.extractors
         )
 
-    def _create_record_counting_extractor(
-        self,
-        extractor_model: Optional[
-            Union[CustomRecordExtractorModel, DpathExtractorModel, CombinedExtractorModel]
-        ],
-        config: Config,
-        decoder: Optional[Decoder],
-    ) -> Optional[RecordExtractor]:
-        """Build the copy of the extractor an `OffsetIncrement` or `PageIncrement` counts a page with.
-
-        The paginator stops when the count falls below the page size, so the count has to be the
-        number of records the API returned. A `CombinedExtractor` with `skip_empty_records` drops
-        empty records before they are counted, which would make a full page holding a null look
-        short and silently end pagination, so every `CombinedExtractor` in the copy is told to count
-        the records it drops.
-        """
-        if not extractor_model:
-            return None
-        extractor: RecordExtractor = self._create_component_from_model(
-            model=extractor_model, config=config, decoder=decoder
-        )
-        self._count_dropped_empty_records(extractor)
-        return extractor
-
     @staticmethod
     def _count_dropped_empty_records(extractor: RecordExtractor) -> None:
+        """Make the record selector count the empty records a `CombinedExtractor` drops.
+
+        `OffsetIncrement` and `PageIncrement` stop when that count falls below the page size, so a
+        full page holding a null would otherwise look short and silently end pagination.
+        """
         if not isinstance(extractor, CombinedExtractor):
             return
         extractor.count_dropped_empty_records = True
@@ -3376,13 +3357,6 @@ def create_offset_increment(
 
         self._reject_union_combined_extractor_for_offset_increment(extractor_model)
 
-        # Ideally we would instantiate the runtime extractor from highest most level (in this case the SimpleRetriever)
-        # so that it can be shared by OffSetIncrement and RecordSelector. However, due to how we instantiate the
-        # decoder with various decorators here, but not in create_record_selector, it is simpler to retain existing
-        # behavior by having two separate extractors with identical behavior since they use the same extractor model.
-        # When we have more time to investigate we can look into reusing the same component.
-        extractor = self._create_record_counting_extractor(extractor_model, config, decoder_to_use)
-
         # Pydantic v1 Union type coercion can convert int to string depending on Union order.
         # If page_size is a string that represents an integer (not an interpolation), convert it back.
         page_size = model.page_size
@@ -3393,26 +3367,14 @@ def create_offset_increment(
             page_size=page_size,
             config=config,
             decoder=decoder_to_use,
-            extractor=extractor,
             inject_on_first_request=model.inject_on_first_request or False,
             parameters=model.parameters or {},
         )
 
+    @staticmethod
     def create_page_increment(
-        self,
-        model: PageIncrementModel,
-        config: Config,
-        decoder: Optional[Decoder] = None,
-        extractor_model: Optional[
-            Union[CustomRecordExtractorModel, DpathExtractorModel, CombinedExtractorModel]
-        ] = None,
-        **kwargs: Any,
+        model: PageIncrementModel, config: Config, **kwargs: Any
     ) -> PageIncrement:
-        # Like OffsetIncrement, we instantiate a separate extractor with identical behavior to the
-        # RecordSelector's so the strategy can count the raw records in the response. This ensures
-        # pagination is driven by the API's page size, not the post-filter record count.
-        extractor = self._create_record_counting_extractor(extractor_model, config, decoder)
-
         # Pydantic v1 Union type coercion can convert int to string depending on Union order.
         # If page_size is a string that represents an integer (not an interpolation), convert it back.
         page_size = model.page_size
@@ -3424,7 +3386,6 @@ def create_page_increment(
             config=config,
             start_from_page=model.start_from_page or 0,
             inject_on_first_request=model.inject_on_first_request or False,
-            extractor=extractor,
             parameters=model.parameters or {},
         )
 
@@ -3623,6 +3584,7 @@ def create_record_selector(
         extractor = self._create_component_from_model(
             model=model.extractor, decoder=decoder, config=config
         )
+        self._count_dropped_empty_records(extractor)
         record_filter = (
             self._create_component_from_model(model.record_filter, config=config)
             if model.record_filter
@@ -4808,6 +4770,7 @@ def create_async_retriever(
         def _get_download_retriever(
             requester: Requester, extractor: RecordExtractor, _decoder: Decoder
         ) -> SimpleRetriever:
+            self._count_dropped_empty_records(extractor)
             # We create a record selector for the download retriever
             # with no schema normalization and no transformations, neither record filter
             # as all this occurs in the record_selector of the AsyncRetriever
@@ -4826,11 +4789,7 @@ def _get_download_retriever(
                     decoder=_decoder,
                     config=config,
                     url_base="",
-                    # Only a `CombinedExtractor` is handed over, so the `union` rejection and the
-                    # count of dropped empty records apply to downloads too. Other extractors keep
-                    # the retriever's own count: a paginator's copy would re-read the response,
-                    # which a streaming download decoder has already consumed, whereas a
-                    # `CombinedExtractor` is only ever built over a buffered decoder.
+                    # Only read by the `union` rejection of `OffsetIncrement`.
                     extractor_model=(
                         model.download_extractor
                         if isinstance(model.download_extractor, CombinedExtractorModel)

airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py

@@ -13,6 +13,7 @@
     PaginationDecoderDecorator,
 )
 from airbyte_cdk.sources.declarative.extractors.dpath_extractor import RecordExtractor
+from airbyte_cdk.sources.declarative.extractors.record_selector import extracted_record_count
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
     PaginationStrategy,
@@ -42,12 +43,14 @@ class OffsetIncrement(PaginationStrategy):
 
     Attributes:
         page_size (InterpolatedString): the number of records to request
+        extractor (Optional[RecordExtractor]): counts the records of a page by extracting them
+            again. Leave it unset to use the count of the `RecordSelector` that read the page
     """
 
     config: Config
     page_size: Optional[Union[str, int]]
     parameters: InitVar[Mapping[str, Any]]
-    extractor: Optional[RecordExtractor]
+    extractor: Optional[RecordExtractor] = None
     decoder: Decoder = field(
         default_factory=lambda: PaginationDecoderDecorator(decoder=JsonDecoder(parameters={}))
     )
@@ -77,29 +80,28 @@ def next_page_token(
         page_size_override: Optional[int] = None,
         stream_slice: Optional[StreamSlice] = None,
     ) -> Optional[Any]:
-        decoded_response = next(self.decoder.decode(response))
-
+        # The records the API returned, not the ones emitted: a record filter must neither end the
+        # pagination nor shorten the offset.
         if self.extractor:
-            page_size_from_response = len(list(self.extractor.extract_records(response=response)))
-            # The extractor could return 0 records which is valid, but evaluates to False. Our fallback in other
-            # cases as the best effort option is to use the incoming last_page_size
-            last_page_size = (
-                page_size_from_response if page_size_from_response is not None else last_page_size
-            )
+            last_page_size = len(list(self.extractor.extract_records(response=response)))
+        else:
+            last_page_size = extracted_record_count(response, default=last_page_size)
 
         # Stop paginating when there are fewer records than the page size or the current page has no records.
         # The comparison uses the page size that was actually requested: after a `REDUCE_PAGE_SIZE` reduction, a
         # full page is smaller than the configured page size and comparing against the latter would end the
         # pagination early and skip records.
-        requested_page_size = (
-            page_size_override
-            if page_size_override is not None
-            else (
-                self._page_size.eval(self.config, response=decoded_response)
-                if self._page_size
-                else None
+        if page_size_override is not None:
+            requested_page_size: Optional[int] = page_size_override
+        elif isinstance(self.page_size, int):
+            # A constant needs no response, which would otherwise be decoded a second time.
+            requested_page_size = self.page_size
+        elif self._page_size:
+            requested_page_size = self._page_size.eval(
+                self.config, response=next(self.decoder.decode(response))
             )
-        )
+        else:
+            requested_page_size = None
         if (requested_page_size and last_page_size < requested_page_size) or last_page_size == 0:
             return None
         elif last_page_token_value is None:

airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py

@@ -9,6 +9,7 @@
 
 from airbyte_cdk.models import FailureType
 from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor
+from airbyte_cdk.sources.declarative.extractors.record_selector import extracted_record_count
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
     PaginationStrategy,
@@ -25,6 +26,8 @@ class PageIncrement(PaginationStrategy):
     Attributes:
         page_size (int): the number of records to request
         start_from_page (int): number of the initial page
+        extractor (Optional[RecordExtractor]): counts the records of a page by extracting them
+            again. Leave it unset to use the count of the `RecordSelector` that read the page
     """
 
     config: Config
@@ -73,11 +76,13 @@ def next_page_token(
                 failure_type=FailureType.config_error,
             )
 
+        # The record count is dependent on the records returned from the response which may not always
+        # align with the size of pages emitted. For example, a record filter can reduce the number of
+        # records observed below the page size even though the API returned a full page.
         if self.extractor:
-            # The record count is dependent on the records returned from the response which may not always
-            # align with the size of pages emitted. For example, a record filter can reduce the number of
-            # records observed below the page size even though the API returned a full page.
             last_page_size = len(list(self.extractor.extract_records(response=response)))
+        else:
+            last_page_size = extracted_record_count(response, default=last_page_size)
 
         # Stop paginating when there are fewer records than the page size or the current page has no records
         if (self._page_size and last_page_size < self._page_size) or last_page_size == 0: