diff --git a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py b/airbyte_cdk/sources/declarative/extractors/combined_extractor.py index bb57c421c3..7a0218466e 100644 --- a/airbyte_cdk/sources/declarative/extractors/combined_extractor.py +++ b/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) " diff --git a/airbyte_cdk/sources/declarative/extractors/record_selector.py b/airbyte_cdk/sources/declarative/extractors/record_selector.py index 80d0155e32..7adf23ea5c 100644 --- a/airbyte_cdk/sources/declarative/extractors/record_selector.py +++ b/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, diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index 1c56386883..0de80fa421 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/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) diff --git a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py index a510731a2e..cb8611c4d6 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py +++ b/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: diff --git a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py index 61b5e94f70..8d24311e88 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py +++ b/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: diff --git a/unit_tests/sources/declarative/extractors/test_combined_extractor.py b/unit_tests/sources/declarative/extractors/test_combined_extractor.py index 0048111896..37c0796e30 100644 --- a/unit_tests/sources/declarative/extractors/test_combined_extractor.py +++ b/unit_tests/sources/declarative/extractors/test_combined_extractor.py @@ -444,17 +444,17 @@ def test_skip_empty_records_does_not_warn_when_nothing_is_dropped(caplog): assert [record.message for record in caplog.records if record.levelno == logging.WARNING] == [] -def _counting_copy(extractor: CombinedExtractor) -> CombinedExtractor: - """What the factory hands a record-counting paginator: the same tree, told to count its drops.""" +def _counting_drops(extractor: CombinedExtractor) -> CombinedExtractor: + """What the factory does to a record selector's extractor: every level counts its drops.""" extractor.count_dropped_empty_records = True for sub_extractor in extractor.extractors: if isinstance(sub_extractor, CombinedExtractor): - _counting_copy(sub_extractor) + _counting_drops(sub_extractor) return extractor -def test_the_counting_copy_counts_dropped_empty_records_under_union(): - extractor = _counting_copy( +def test_a_drop_counting_extractor_counts_dropped_empty_records_under_union(): + extractor = _counting_drops( CombinedExtractor( extractors=[ _CountingExtractor(records=[{"id": 1}, None, {}]), @@ -469,7 +469,7 @@ def test_the_counting_copy_counts_dropped_empty_records_under_union(): assert len(list(extractor.extract_records(create_response(GRAPHQL_BODY)))) == 4 -def test_the_counting_copy_counts_the_same_winner_the_record_stream_reads(): +def test_a_drop_counting_extractor_counts_the_same_winner_the_record_stream_reads(): """A page of nulls must not win the count while the record stream falls through past it. The winner is chosen by the surviving records, so the count is the fallback path's, dropped @@ -488,7 +488,7 @@ def test_the_counting_copy_counts_the_same_winner_the_record_stream_reads(): assert ( len( list( - _counting_copy(_monday_extractor(skip_empty_records=True)).extract_records( + _counting_drops(_monday_extractor(skip_empty_records=True)).extract_records( create_response(body) ) ) @@ -497,8 +497,8 @@ def test_the_counting_copy_counts_the_same_winner_the_record_stream_reads(): ) -def test_the_counting_copy_counts_nothing_when_every_record_is_empty(): - extractor = _counting_copy( +def test_a_drop_counting_extractor_counts_nothing_when_every_record_is_empty(): + extractor = _counting_drops( CombinedExtractor( extractors=[_CountingExtractor(records=[None]), _CountingExtractor(records=[{}])], mode=CombineMode.first_match, @@ -510,8 +510,8 @@ def test_the_counting_copy_counts_nothing_when_every_record_is_empty(): assert list(extractor.extract_records(create_response(GRAPHQL_BODY))) == [] -def test_the_counting_copy_counts_the_drops_of_a_nested_combined_extractor(): - extractor = _counting_copy( +def test_a_drop_counting_extractor_counts_the_drops_of_a_nested_combined_extractor(): + extractor = _counting_drops( CombinedExtractor( extractors=[ CombinedExtractor( @@ -535,9 +535,9 @@ def test_the_counting_copy_counts_the_drops_of_a_nested_combined_extractor(): assert len(list(extractor.extract_records(create_response(GRAPHQL_BODY)))) == 2 -def test_the_counting_copy_does_not_warn(caplog): - """The record stream already warned about the page the paginator re-reads.""" - extractor = _counting_copy( +def test_a_drop_counting_extractor_warns_about_the_records_it_drops(caplog): + """It is the only extractor reading the page, so nothing else reports the drop.""" + extractor = _counting_drops( CombinedExtractor( extractors=[_CountingExtractor(records=[{"id": 1}, None])], skip_empty_records=True, @@ -548,7 +548,9 @@ def test_the_counting_copy_does_not_warn(caplog): with caplog.at_level(logging.WARNING, logger="airbyte"): list(extractor.extract_records(create_response(GRAPHQL_BODY))) - assert not [record for record in caplog.records if record.levelno == logging.WARNING] + warnings = [record.message for record in caplog.records if record.levelno == logging.WARNING] + assert len(warnings) == 1 + assert "dropped 1 empty record(s)" in warnings[0] def test_the_documented_first_match_example_yields_one_record_per_item(): diff --git a/unit_tests/sources/declarative/extractors/test_record_selector.py b/unit_tests/sources/declarative/extractors/test_record_selector.py index 2d32bc75cc..ed256a1582 100644 --- a/unit_tests/sources/declarative/extractors/test_record_selector.py +++ b/unit_tests/sources/declarative/extractors/test_record_selector.py @@ -11,7 +11,10 @@ from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder from airbyte_cdk.sources.declarative.extractors.dpath_extractor import DpathExtractor from airbyte_cdk.sources.declarative.extractors.record_filter import RecordFilter -from airbyte_cdk.sources.declarative.extractors.record_selector import RecordSelector +from airbyte_cdk.sources.declarative.extractors.record_selector import ( + RecordSelector, + extracted_record_count, +) from airbyte_cdk.sources.declarative.transformations import RecordTransformation from airbyte_cdk.sources.types import Record, StreamSlice from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer @@ -312,3 +315,48 @@ def transformation_side_effect(record, config, stream_state, stream_slice): else: assert final_record_data[0]["id"] == 2 assert final_record_data[0]["myfield"] == 999 + + +class _StopReadingAtTheFirstRejectedRecord(RecordFilter): + def filter_records(self, records, stream_state, stream_slice=None, next_page_token=None): + for record in records: + if record["id"] == 2: + return + yield record + + +@pytest.mark.parametrize( + "record_filter, expected_ids", + [ + pytest.param( + RecordFilter(config={}, condition="{{ record['id'] != 2 }}", parameters={}), + [1, 3], + id="record_filter", + ), + pytest.param( + _StopReadingAtTheFirstRejectedRecord(config={}, parameters={}), + [1], + id="filter_that_stops_reading_the_page", + ), + ], +) +def test_select_records_counts_the_records_of_the_page_before_filtering_them( + record_filter, expected_ids +): + record_selector = RecordSelector( + extractor=DpathExtractor(field_path=["data"], config={}, parameters={}), + record_filter=record_filter, + config={}, + name="test_stream", + schema_normalization=TypeTransformer(TransformConfig.NoTransform), + parameters={}, + ) + response = create_response({"data": [{"id": 1}, {"id": 2}, {"id": 3}]}) + assert extracted_record_count(response, default=-1) == -1 + + records = list( + record_selector.select_records(response=response, stream_state={}, records_schema={}) + ) + + assert [record.data["id"] for record in records] == expected_ids + assert extracted_record_count(response, default=-1) == 3 diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index cc16d3a975..3e3b3e4ac1 100644 --- a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +++ b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py @@ -4427,15 +4427,12 @@ def test_create_offset_increment(): inject_on_first_request=True, ) - expected_extractor = DpathExtractor(field_path=["results"], config=input_config, parameters={}) - extractor_model = DpathExtractorModel( - type="DpathExtractor", field_path=expected_extractor.field_path - ) + extractor_model = DpathExtractorModel(type="DpathExtractor", field_path=["results"]) expected_strategy = OffsetIncrement( page_size=10, inject_on_first_request=True, - extractor=expected_extractor, + extractor=None, parameters={}, config=input_config, ) @@ -4448,8 +4445,7 @@ def test_create_offset_increment(): assert strategy.inject_on_first_request == expected_strategy.inject_on_first_request assert strategy.config == input_config - assert isinstance(strategy.extractor, DpathExtractor) - assert strategy.extractor.field_path == expected_extractor.field_path + assert strategy.extractor is None class MyCustomSchemaLoader(SchemaLoader): @@ -9691,12 +9687,10 @@ def test_first_match_combined_extractor_counts_only_the_winning_sub_extractor_fo assert isinstance(retriever.record_selector.extractor, CombinedExtractor) pagination_strategy = retriever.paginator.pagination_strategy assert isinstance(pagination_strategy, OffsetIncrement) - # The paginator counts records with its own copy, so the response is traversed twice per page. - assert isinstance(pagination_strategy.extractor, CombinedExtractor) - assert pagination_strategy.extractor is not retriever.record_selector.extractor # `a` misses, so the winner is `b` with its two records: the offset advances by 2, not by 4. response = _json_response({"a": [], "b": [{"id": 1}, {"id": 2}]}) + list(retriever.record_selector.select_records(response, {}, {})) assert ( pagination_strategy.next_page_token( response=response, last_page_size=2, last_record=None, last_page_token_value=0 @@ -9717,7 +9711,6 @@ def test_union_combined_extractor_is_accepted_by_a_page_increment_paginator(): pagination_strategy = retriever.paginator.pagination_strategy assert isinstance(pagination_strategy, PageIncrement) - assert isinstance(pagination_strategy.extractor, CombinedExtractor) _FIRST_MATCH_SKIPPING_EMPTY_RECORDS = """ type: CombinedExtractor @@ -9745,16 +9738,12 @@ def test_a_record_counting_paginator_counts_the_empty_records_the_extractor_drop _retriever_manifest_with_paginator(_FIRST_MATCH_SKIPPING_EMPTY_RECORDS, pagination_strategy) ) - record_selector_extractor = retriever.record_selector.extractor - counting_extractor = retriever.paginator.pagination_strategy.extractor - assert isinstance(record_selector_extractor, CombinedExtractor) - assert isinstance(counting_extractor, CombinedExtractor) - assert record_selector_extractor.count_dropped_empty_records is False - assert counting_extractor.count_dropped_empty_records is True - full_page_with_a_null = _json_response({"a": [None, {"id": 1}]}) - assert list(record_selector_extractor.extract_records(full_page_with_a_null)) == [{"id": 1}] - # `last_page_size` is what the retriever emitted; the strategy must count the page itself. + assert [ + record.data + for record in retriever.record_selector.select_records(full_page_with_a_null, {}, {}) + ] == [{"id": 1}] + # `last_page_size` is what the retriever emitted; the strategy must count the whole page. assert ( retriever.paginator.pagination_strategy.next_page_token( response=full_page_with_a_null, @@ -9766,6 +9755,85 @@ def test_a_record_counting_paginator_counts_the_empty_records_the_extractor_drop ) +def _filtering_retriever( + pagination_strategy_type: str, decoder: Optional[Mapping[str, Any]] = None +) -> SimpleRetriever: + """Requests pages of two records and drops the record with id 2.""" + return _build_retriever( + { + "type": "SimpleRetriever", + **({"decoder": decoder} if decoder else {}), + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "/records", + "http_method": "GET", + }, + "paginator": { + "type": "DefaultPaginator", + "pagination_strategy": {"type": pagination_strategy_type, "page_size": 2}, + "page_token_option": { + "type": "RequestOption", + "inject_into": "request_parameter", + "field_name": "token", + }, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": ["results"]}, + "record_filter": {"type": "RecordFilter", "condition": "{{ record['id'] != 2 }}"}, + }, + } + ) + + +@pytest.mark.parametrize( + "pagination_strategy_type, second_page_token", + [("OffsetIncrement", "2"), ("PageIncrement", "1")], +) +def test_a_record_counting_paginator_extracts_each_page_once( + requests_mock, mocker, pagination_strategy_type: str, second_page_token: str +): + """The filter drops a record of the full first page, which must not end the pagination.""" + retriever = _filtering_retriever(pagination_strategy_type) + requests_mock.get( + "https://api.test.com/records", + [{"json": {"results": [{"id": 1}, {"id": 2}]}}, {"json": {"results": [{"id": 3}]}}], + ) + extract_records = mocker.spy(DpathExtractor, "extract_records") + + records = list(retriever.read_records({}, StreamSlice(partition={}, cursor_slice={}))) + + assert [record.data["id"] for record in records] == [1, 3] + assert [request.qs for request in requests_mock.request_history] == [ + {}, + {"token": [second_page_token]}, + ] + assert extract_records.call_count == 2 + + +@pytest.mark.parametrize("pagination_strategy_type", ["OffsetIncrement", "PageIncrement"]) +def test_a_record_counting_paginator_paginates_a_streamed_response( + requests_mock, pagination_strategy_type: str +): + """A streamed body can only be read once: counting the page by reading it again found none.""" + retriever = _filtering_retriever( + pagination_strategy_type, + decoder={"type": "GzipDecoder", "decoder": {"type": "JsonDecoder"}}, + ) + requests_mock.get( + "https://api.test.com/records", + [ + {"content": gzip.compress(json.dumps({"results": page}).encode())} + for page in ([{"id": 1}, {"id": 3}], [{"id": 4}]) + ], + ) + + records = list(retriever.read_records({}, StreamSlice(partition={}, cursor_slice={}))) + + assert [record.data["id"] for record in records] == [1, 3, 4] + + def _async_retriever_definition( download_extractor: Mapping[str, Any], download_paginator: Mapping[str, Any] ) -> Dict[str, Any]: @@ -9818,7 +9886,7 @@ def _download_paginator(pagination_strategy_type: str) -> Dict[str, Any]: } -def _build_download_paginator(definition: Mapping[str, Any]) -> Any: +def _build_download_retriever(definition: Mapping[str, Any]) -> Any: component = factory.create_component( model_type=AsyncRetrieverModel, component_definition=definition, @@ -9831,7 +9899,7 @@ def _build_download_paginator(definition: Mapping[str, Any]) -> Any: job_repository = component.stream_slicer.job_orchestrator_factory( [StreamSlice(partition={}, cursor_slice={})] )._job_repository - return job_repository.download_retriever.paginator + return job_repository.download_retriever def test_union_combined_download_extractor_is_rejected_by_an_offset_increment_download_paginator(): @@ -9853,7 +9921,7 @@ def test_union_combined_download_extractor_is_rejected_by_an_offset_increment_do ) with pytest.raises(AirbyteTracedException) as exc_info: - _build_download_paginator(definition) + _build_download_retriever(definition) assert exc_info.value.failure_type == FailureType.config_error assert "union" in exc_info.value.message @@ -9874,13 +9942,14 @@ def test_combined_download_extractor_skipping_empty_records_counts_them_for_the_ _download_paginator("PageIncrement"), ) - pagination_strategy = _build_download_paginator(definition).pagination_strategy + download_retriever = _build_download_retriever(definition) + pagination_strategy = download_retriever.paginator.pagination_strategy assert isinstance(pagination_strategy, PageIncrement) - assert isinstance(pagination_strategy.extractor, CombinedExtractor) - assert pagination_strategy.extractor.count_dropped_empty_records is True + response = _json_response({"a": [None, {"id": 1}]}) + list(download_retriever.record_selector.select_records(response, {}, {})) assert ( pagination_strategy.next_page_token( - response=_json_response({"a": [None, {"id": 1}]}), + response=response, last_page_size=1, last_record={"id": 1}, last_page_token_value=0, @@ -9896,7 +9965,7 @@ def test_download_paginator_keeps_the_retriever_count_for_other_download_extract _download_paginator("OffsetIncrement"), ) - pagination_strategy = _build_download_paginator(definition).pagination_strategy + pagination_strategy = _build_download_retriever(definition).paginator.pagination_strategy assert isinstance(pagination_strategy, OffsetIncrement) assert pagination_strategy.extractor is None diff --git a/unit_tests/sources/declarative/requesters/paginators/test_offset_increment.py b/unit_tests/sources/declarative/requesters/paginators/test_offset_increment.py index a12bd1e5a1..bacbcbeff7 100644 --- a/unit_tests/sources/declarative/requesters/paginators/test_offset_increment.py +++ b/unit_tests/sources/declarative/requesters/paginators/test_offset_increment.py @@ -4,14 +4,17 @@ import json from typing import Any, Optional +from unittest.mock import Mock import pytest import requests -from airbyte_cdk.sources.declarative.extractors import DpathExtractor +from airbyte_cdk.sources.declarative.decoders import JsonDecoder +from airbyte_cdk.sources.declarative.extractors import DpathExtractor, RecordFilter, RecordSelector from airbyte_cdk.sources.declarative.requesters.paginators.strategies.offset_increment import ( OffsetIncrement, ) +from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer @pytest.mark.parametrize( @@ -241,3 +244,31 @@ def test_given_page_size_interpolates_to_empty_string_then_paginate_until_an_emp ) is None ) + + +def test_given_a_page_counted_by_the_record_selector_then_do_not_read_the_response_again(): + """The filter dropped the whole page, yet the offset advances by the two records returned.""" + response = _response([{"id": 1}, {"id": 2}]) + record_selector = RecordSelector( + extractor=DpathExtractor(field_path=["results"], config={}, parameters={}), + record_filter=RecordFilter(config={}, condition="{{ False }}", parameters={}), + config={}, + name="test_stream", + schema_normalization=TypeTransformer(TransformConfig.NoTransform), + parameters={}, + ) + assert not list( + record_selector.select_records(response=response, stream_state={}, records_schema={}) + ) + decoder = Mock(wraps=JsonDecoder(parameters={})) + strategy = OffsetIncrement( + page_size=2, extractor=None, decoder=decoder, config={}, parameters={} + ) + + assert ( + strategy.next_page_token( + response=response, last_page_size=0, last_record=None, last_page_token_value=4 + ) + == 6 + ) + decoder.decode.assert_not_called() diff --git a/unit_tests/sources/declarative/requesters/paginators/test_page_increment.py b/unit_tests/sources/declarative/requesters/paginators/test_page_increment.py index ef8339fe55..e688585098 100644 --- a/unit_tests/sources/declarative/requesters/paginators/test_page_increment.py +++ b/unit_tests/sources/declarative/requesters/paginators/test_page_increment.py @@ -9,10 +9,11 @@ import requests from airbyte_cdk.models import FailureType -from airbyte_cdk.sources.declarative.extractors import DpathExtractor +from airbyte_cdk.sources.declarative.extractors import DpathExtractor, RecordFilter, RecordSelector from airbyte_cdk.sources.declarative.requesters.paginators.strategies.page_increment import ( PageIncrement, ) +from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer from airbyte_cdk.utils.traced_exception import AirbyteTracedException @@ -183,3 +184,28 @@ def test_given_page_size_override_then_raise_config_error(): ) assert exception.value.failure_type == FailureType.config_error + + +def test_given_a_page_counted_by_the_record_selector_then_use_its_count(): + """The filter dropped the whole page, yet the page is full, so the pagination goes on.""" + response = requests.Response() + response._content = json.dumps({"results": [{"id": 1}, {"id": 2}]}).encode("utf-8") + record_selector = RecordSelector( + extractor=DpathExtractor(field_path=["results"], config={}, parameters={}), + record_filter=RecordFilter(config={}, condition="{{ False }}", parameters={}), + config={}, + name="test_stream", + schema_normalization=TypeTransformer(TransformConfig.NoTransform), + parameters={}, + ) + assert not list( + record_selector.select_records(response=response, stream_state={}, records_schema={}) + ) + strategy = PageIncrement(page_size=2, start_from_page=1, config={}, parameters={}) + + assert ( + strategy.next_page_token( + response=response, last_page_size=0, last_record=None, last_page_token_value=3 + ) + == 4 + )