Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -3032,10 +3032,23 @@ def create_offset_increment(
parameters=model.parameters or {},
)

@staticmethod
def create_page_increment(
model: PageIncrementModel, config: Config, **kwargs: Any
self,
model: PageIncrementModel,
config: Config,
decoder: Optional[Decoder] = None,
extractor_model: Optional[Union[CustomRecordExtractorModel, DpathExtractorModel]] = None,
**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_component_from_model(model=extractor_model, config=config, decoder=decoder)
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
Expand All @@ -3047,6 +3060,7 @@ 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 {},
)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import requests

from airbyte_cdk.sources.declarative.extractors.record_extractor import RecordExtractor
from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
from airbyte_cdk.sources.declarative.requesters.paginators.strategies.pagination_strategy import (
PaginationStrategy,
Expand All @@ -29,6 +30,7 @@ class PageIncrement(PaginationStrategy):
parameters: InitVar[Mapping[str, Any]]
start_from_page: int = 0
inject_on_first_request: bool = False
extractor: Optional[RecordExtractor] = None

def __post_init__(self, parameters: Mapping[str, Any]) -> None:
if isinstance(self.page_size, int) or (self.page_size is None):
Expand All @@ -52,6 +54,12 @@ def next_page_token(
last_record: Optional[Record],
last_page_token_value: Optional[Any],
) -> Optional[Any]:
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)))

# 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:
return None
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,66 @@
import pytest
import requests

from airbyte_cdk.sources.declarative.extractors import DpathExtractor
from airbyte_cdk.sources.declarative.requesters.paginators.strategies.page_increment import (
PageIncrement,
)


@pytest.mark.parametrize(
"response_results, last_page_size, last_page_token_value, expected_next_page_token",
[
pytest.param(
[{"id": 1}, {"id": 2}],
0,
3,
4,
id="test_full_page_continues_even_if_all_records_filtered",
),
pytest.param(
[{"id": 1}, {"id": 2}],
1,
3,
4,
id="test_full_page_continues_even_if_some_records_filtered",
),
pytest.param(
[{"id": 1}],
1,
3,
None,
id="test_partial_page_stops_pagination",
),
pytest.param(
[],
0,
3,
None,
id="test_empty_page_stops_pagination",
),
],
)
def test_page_increment_paginator_strategy_with_extractor(
response_results, last_page_size, last_page_token_value, expected_next_page_token
):
extractor = DpathExtractor(field_path=["results"], parameters={}, config={})
paginator_strategy = PageIncrement(
page_size=2, parameters={}, start_from_page=1, extractor=extractor, config={}
)

response = requests.Response()
response.headers = {"A_HEADER": "HEADER_VALUE"}
response._content = json.dumps({"results": response_results}).encode("utf-8")

next_page_token = paginator_strategy.next_page_token(
response,
last_page_size,
response_results[-1] if response_results else None,
last_page_token_value,
)
assert expected_next_page_token == next_page_token


@pytest.mark.parametrize(
"page_size, start_from, last_page_size, last_record, last_page_token_value, expected_next_page_token, expected_offset",
[
Expand Down
Loading