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 c2551b485..cfc01b8e3 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1771,7 +1771,12 @@ def _validate_jitter_range(jitter_range_in_seconds: Optional[float]) -> None: raise ValueError("jitter_range_in_seconds must be greater than or equal to 0") def create_cursor_pagination( - self, model: CursorPaginationModel, config: Config, decoder: Decoder, **kwargs: Any + self, + model: CursorPaginationModel, + config: Config, + decoder: Decoder, + extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None, + **kwargs: Any, ) -> CursorPaginationStrategy: if isinstance(decoder, PaginationDecoderDecorator): inner_decoder = decoder.decoder @@ -1786,6 +1791,17 @@ def create_cursor_pagination( self._UNSUPPORTED_DECODER_ERROR.format(decoder_type=type(inner_decoder)) ) + # Like OffsetIncrement and PageIncrement, we instantiate a separate extractor with identical + # behavior to the RecordSelector's so the strategy can count the raw records in the response, + # keeping `last_page_size` in stop_condition/cursor_value interpolation pre-record-filter. + extractor = ( + self._create_component_from_model( + model=extractor_model, config=config, decoder=decoder_to_use + ) + if extractor_model + else None + ) + # 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 @@ -1795,6 +1811,7 @@ def create_cursor_pagination( return CursorPaginationStrategy( cursor_value=model.cursor_value, decoder=decoder_to_use, + extractor=extractor, page_size=page_size, stop_condition=model.stop_condition, config=config, diff --git a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/cursor_pagination_strategy.py b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/cursor_pagination_strategy.py index 73a644b5b..1909202d2 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/cursor_pagination_strategy.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/cursor_pagination_strategy.py @@ -12,6 +12,7 @@ JsonDecoder, PaginationDecoderDecorator, ) +from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import ( @@ -41,6 +42,7 @@ class CursorPaginationStrategy(PaginationStrategy): decoder: Decoder = field( default_factory=lambda: PaginationDecoderDecorator(decoder=JsonDecoder(parameters={})) ) + extractor: Optional[RecordExtractor] = None def __post_init__(self, parameters: Mapping[str, Any]) -> None: if isinstance(self.cursor_value, str): @@ -83,6 +85,13 @@ def next_page_token( # is not indexable or useful for parsing the cursor, so we replace it with the link dictionary from response.links headers: Dict[str, Any] = dict(response.headers) headers["link"] = response.links + + if self.extractor: + # The incoming last_page_size is the number of records emitted after the RecordSelector ran, so a + # record filter (e.g. client-side incremental) can make it smaller than the page the API returned. + # Stop conditions like `{{ last_page_size < page_size }}` would then stop pagination prematurely. + # Re-counting the raw records in the response keeps pagination driven by the API's page size. + last_page_size = len(list(self.extractor.extract_records(response=response))) if self._stop_condition: should_stop = self._stop_condition.eval( self.config, 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 5149f46d2..a698ad373 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 @@ -63,6 +63,7 @@ CompositeErrorHandler as CompositeErrorHandlerModel, ) from airbyte_cdk.sources.declarative.models import ConcurrencyLevel as ConcurrencyLevelModel +from airbyte_cdk.sources.declarative.models import CursorPagination as CursorPaginationModel from airbyte_cdk.sources.declarative.models import CustomErrorHandler as CustomErrorHandlerModel from airbyte_cdk.sources.declarative.models import ( CustomPartitionRouter as CustomPartitionRouterModel, @@ -2318,7 +2319,7 @@ def test_create_default_paginator(): component_definition=paginator_manifest, config=input_config, url_base="https://airbyte.io", - extractor_model=DpathExtractor(field_path=["results"], config=input_config, parameters={}), + extractor_model=DpathExtractorModel(type="DpathExtractor", field_path=["results"]), decoder=JsonDecoder(parameters={}), ) @@ -2328,6 +2329,8 @@ def test_create_default_paginator(): assert isinstance(paginator.pagination_strategy, CursorPaginationStrategy) assert paginator.pagination_strategy.page_size == 50 assert paginator.pagination_strategy._cursor_value.string == "{{ response._metadata.next }}" + assert isinstance(paginator.pagination_strategy.extractor, DpathExtractor) + assert paginator.pagination_strategy.extractor.field_path == ["results"] assert isinstance(paginator.page_size_option, RequestOption) assert paginator.page_size_option.inject_into == RequestOptionType.request_parameter @@ -3185,6 +3188,30 @@ def test_create_offset_increment(): assert strategy.extractor.field_path == expected_extractor.field_path +def test_create_cursor_pagination(): + model = CursorPaginationModel( + type="CursorPagination", + cursor_value="{{ response.next_token }}", + stop_condition="{{ last_page_size < 2 }}", + page_size=2, + ) + + expected_extractor = DpathExtractor(field_path=["results"], config=input_config, parameters={}) + extractor_model = DpathExtractorModel( + type="DpathExtractor", field_path=expected_extractor.field_path + ) + + strategy = factory.create_cursor_pagination( + model, input_config, extractor_model=extractor_model, decoder=JsonDecoder(parameters={}) + ) + + assert strategy.get_page_size() == 2 + assert strategy.config == input_config + + assert isinstance(strategy.extractor, DpathExtractor) + assert strategy.extractor.field_path == expected_extractor.field_path + + class MyCustomSchemaLoader(SchemaLoader): def get_json_schema(self) -> Mapping[str, Any]: """Returns a mapping describing the stream's schema""" diff --git a/unit_tests/sources/declarative/requesters/paginators/test_cursor_pagination_strategy.py b/unit_tests/sources/declarative/requesters/paginators/test_cursor_pagination_strategy.py index da21e1074..a0446abb3 100644 --- a/unit_tests/sources/declarative/requesters/paginators/test_cursor_pagination_strategy.py +++ b/unit_tests/sources/declarative/requesters/paginators/test_cursor_pagination_strategy.py @@ -8,6 +8,7 @@ import requests from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder +from airbyte_cdk.sources.declarative.extractors import DpathExtractor from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean from airbyte_cdk.sources.declarative.requesters.paginators.strategies.cursor_pagination_strategy import ( CursorPaginationStrategy, @@ -87,6 +88,77 @@ def test_cursor_pagination_strategy(template_string, stop_condition, expected_to assert page_size == strategy.get_page_size() +@pytest.mark.parametrize( + "response_results, last_page_size, expected_next_page_token", + [ + pytest.param( + [{"id": 1}, {"id": 2}], + 0, + "next_token", + id="test_full_page_continues_even_if_all_records_filtered", + ), + pytest.param( + [{"id": 1}, {"id": 2}], + 1, + "next_token", + id="test_full_page_continues_even_if_some_records_filtered", + ), + pytest.param( + [{"id": 1}], + 1, + None, + id="test_partial_page_stops_pagination", + ), + pytest.param( + [], + 0, + None, + id="test_empty_page_stops_pagination", + ), + ], +) +def test_cursor_pagination_strategy_with_extractor( + response_results, last_page_size, expected_next_page_token +): + extractor = DpathExtractor(field_path=["results"], parameters={}, config={}) + strategy = CursorPaginationStrategy( + page_size=2, + cursor_value="{{ response.next_token }}", + stop_condition="{{ last_page_size < 2 }}", + extractor=extractor, + config={}, + parameters={}, + ) + + response = requests.Response() + response._content = json.dumps( + {"results": response_results, "next_token": "next_token"} + ).encode("utf-8") + last_record = ( + Record(data=response_results[-1], stream_name="stream_name") if response_results else None + ) + + next_page_token = strategy.next_page_token(response, last_page_size, last_record) + assert expected_next_page_token == next_page_token + + +def test_cursor_pagination_strategy_with_extractor_interpolates_raw_count_in_cursor_value(): + extractor = DpathExtractor(field_path=["results"], parameters={}, config={}) + strategy = CursorPaginationStrategy( + page_size=2, + cursor_value="{{ last_page_size }}", + extractor=extractor, + config={}, + parameters={}, + ) + + response = requests.Response() + response._content = json.dumps({"results": [{"id": 1}, {"id": 2}]}).encode("utf-8") + + next_page_token = strategy.next_page_token(response, 0, None) + assert next_page_token == 2 + + def test_last_record_points_to_the_last_item_in_last_records_array(): last_records = [{"id": 0, "more_records": True}, {"id": 1, "more_records": True}] strategy = CursorPaginationStrategy(