diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 7d99a01881..abb16a7d88 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -2421,6 +2421,7 @@ definitions: - RESET_PAGINATION - RATE_LIMITED - REFRESH_TOKEN_THEN_RETRY + - REDUCE_PAGE_SIZE examples: - SUCCESS - FAIL @@ -2429,6 +2430,7 @@ definitions: - RESET_PAGINATION - RATE_LIMITED - REFRESH_TOKEN_THEN_RETRY + - REDUCE_PAGE_SIZE failure_type: title: Failure Type description: Failure type of traced exception if a response matches the filter. diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 3bee30ca68..ae8b09c0c6 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -557,6 +557,7 @@ class Action(Enum): RESET_PAGINATION = "RESET_PAGINATION" RATE_LIMITED = "RATE_LIMITED" REFRESH_TOKEN_THEN_RETRY = "REFRESH_TOKEN_THEN_RETRY" + REDUCE_PAGE_SIZE = "REDUCE_PAGE_SIZE" class FailureType(Enum): @@ -578,6 +579,7 @@ class HttpResponseFilter(BaseModel): "RESET_PAGINATION", "RATE_LIMITED", "REFRESH_TOKEN_THEN_RETRY", + "REDUCE_PAGE_SIZE", ], title="Action", ) 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 73a644b5ba..6457b69812 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 @@ -61,6 +61,7 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: if not isinstance(page_size, int): raise Exception(f"{page_size} is of type {type(page_size)}. Expected {int}") self._page_size = page_size + self._default_page_size = self._page_size @property def initial_token(self) -> Optional[Any]: @@ -104,3 +105,10 @@ def next_page_token( def get_page_size(self) -> Optional[int]: return self._page_size + + def reduce_page_size(self) -> None: + if self._page_size is not None and self._page_size > 1: + self._page_size = max(1, self._page_size // 2) + + def reset_page_size(self) -> None: + self._page_size = self._default_page_size 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 4370155dec..f8e2402de2 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py @@ -61,6 +61,8 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: ) else: self._page_size = None + self._default_page_size = self._page_size + self._effective_page_size: Optional[int] = None @property def initial_token(self) -> Optional[Any]: @@ -103,6 +105,8 @@ def next_page_token( return last_page_token_value + last_page_size def get_page_size(self) -> Optional[int]: + if self._effective_page_size is not None: + return self._effective_page_size if self._page_size: page_size = self._page_size.eval(self.config) if not isinstance(page_size, int): @@ -110,3 +114,19 @@ def get_page_size(self) -> Optional[int]: return page_size else: return None + + def _get_default_page_size(self) -> Optional[int]: + if self._default_page_size: + page_size = self._default_page_size.eval(self.config) + if not isinstance(page_size, int): + raise Exception(f"{page_size} is of type {type(page_size)}. Expected {int}") + return page_size + return None + + def reduce_page_size(self) -> None: + current = self.get_page_size() + if current is not None and current > 1: + self._effective_page_size = max(1, current // 2) + + def reset_page_size(self) -> None: + self._effective_page_size = 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 2e1643b565..62e58e4aad 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/page_increment.py @@ -38,6 +38,7 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: if not isinstance(page_size, int): raise Exception(f"{page_size} is of type {type(page_size)}. Expected {int}") self._page_size = page_size + self._default_page_size = self._page_size @property def initial_token(self) -> Optional[Any]: @@ -69,3 +70,10 @@ def next_page_token( def get_page_size(self) -> Optional[int]: return self._page_size + + def reduce_page_size(self) -> None: + if self._page_size is not None and self._page_size > 1: + self._page_size = max(1, self._page_size // 2) + + def reset_page_size(self) -> None: + self._page_size = self._default_page_size diff --git a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py index dae02ba138..85780ccfb2 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/pagination_strategy.py @@ -46,3 +46,17 @@ def get_page_size(self) -> Optional[int]: """ :return: page size: The number of records to fetch in a page. Returns None if unspecified """ + + def reduce_page_size(self) -> None: + """Halve the current effective page size (floored at 1). + + Called by `SimpleRetriever` when a `REDUCE_PAGE_SIZE` response action is received. + Subclasses that support dynamic page-size reduction should override this method. + """ + + def reset_page_size(self) -> None: + """Restore the page size to the originally configured default. + + Called by `SimpleRetriever` after a successful page fetch following a reduction. + Subclasses that support dynamic page-size reduction should override this method. + """ diff --git a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py index 1f2eb1c668..dc45cdc34d 100644 --- a/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py +++ b/airbyte_cdk/sources/declarative/retrievers/simple_retriever.py @@ -28,6 +28,9 @@ from airbyte_cdk.sources.declarative.partition_routers.single_partition_router import ( SinglePartitionRouter, ) +from airbyte_cdk.sources.declarative.requesters.paginators.default_paginator import ( + DefaultPaginator, +) from airbyte_cdk.sources.declarative.requesters.paginators.no_pagination import NoPagination from airbyte_cdk.sources.declarative.requesters.paginators.paginator import Paginator from airbyte_cdk.sources.declarative.requesters.query_properties import QueryProperties @@ -41,6 +44,9 @@ from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer from airbyte_cdk.sources.source import ExperimentalClassWarning from airbyte_cdk.sources.streams.core import StreamData +from airbyte_cdk.sources.streams.http.page_size_reduction_exception import ( + PageSizeReductionRequiredException, +) from airbyte_cdk.sources.streams.http.pagination_reset_exception import ( PaginationResetRequiredException, ) @@ -403,7 +409,11 @@ def _read_pages( yield current_record except PaginationResetRequiredException: reset_pagination = True + except PageSizeReductionRequiredException: + self._reduce_paginator_page_size() + continue else: + self._reset_paginator_page_size() if not response: break @@ -433,6 +443,25 @@ def _read_pages( # Always return an empty generator just in case no records were ever yielded yield from [] + def _reduce_paginator_page_size(self) -> None: + """Delegate page-size reduction to the paginator's `PaginationStrategy`, if available.""" + if isinstance(self._paginator, DefaultPaginator): + strategy = self._paginator.pagination_strategy + previous = strategy.get_page_size() + strategy.reduce_page_size() + current = strategy.get_page_size() + LOGGER.info( + "Reducing page size for stream '%s' from %s to %s due to server error.", + self.name, + previous, + current, + ) + + def _reset_paginator_page_size(self) -> None: + """Restore the paginator's page size to the configured default after a successful fetch.""" + if isinstance(self._paginator, DefaultPaginator): + self._paginator.pagination_strategy.reset_page_size() + def _get_initial_next_page_token(self) -> Optional[Mapping[str, Any]]: initial_token = self._paginator.get_initial_token() next_page_token = {"next_page_token": initial_token} if initial_token is not None else None diff --git a/airbyte_cdk/sources/streams/http/error_handlers/response_models.py b/airbyte_cdk/sources/streams/http/error_handlers/response_models.py index 082d580d53..10416c1a7e 100644 --- a/airbyte_cdk/sources/streams/http/error_handlers/response_models.py +++ b/airbyte_cdk/sources/streams/http/error_handlers/response_models.py @@ -19,6 +19,7 @@ class ResponseAction(Enum): RESET_PAGINATION = "RESET_PAGINATION" RATE_LIMITED = "RATE_LIMITED" REFRESH_TOKEN_THEN_RETRY = "REFRESH_TOKEN_THEN_RETRY" + REDUCE_PAGE_SIZE = "REDUCE_PAGE_SIZE" @dataclass diff --git a/airbyte_cdk/sources/streams/http/http_client.py b/airbyte_cdk/sources/streams/http/http_client.py index c1d0eabd67..c806150a85 100644 --- a/airbyte_cdk/sources/streams/http/http_client.py +++ b/airbyte_cdk/sources/streams/http/http_client.py @@ -42,6 +42,9 @@ RequestBodyException, UserDefinedBackoffException, ) +from airbyte_cdk.sources.streams.http.page_size_reduction_exception import ( + PageSizeReductionRequiredException, +) from airbyte_cdk.sources.streams.http.pagination_reset_exception import ( PaginationResetRequiredException, ) @@ -441,6 +444,9 @@ def _handle_error_resolution( if error_resolution.response_action == ResponseAction.RESET_PAGINATION: raise PaginationResetRequiredException() + if error_resolution.response_action == ResponseAction.REDUCE_PAGE_SIZE: + raise PageSizeReductionRequiredException() + # Emit stream status RUNNING with the reason RATE_LIMITED to log that the rate limit has been reached if error_resolution.response_action == ResponseAction.RATE_LIMITED: # TODO: Update to handle with message repository when concurrent message repository is ready diff --git a/airbyte_cdk/sources/streams/http/page_size_reduction_exception.py b/airbyte_cdk/sources/streams/http/page_size_reduction_exception.py new file mode 100644 index 0000000000..0ee658ff5b --- /dev/null +++ b/airbyte_cdk/sources/streams/http/page_size_reduction_exception.py @@ -0,0 +1,2 @@ +class PageSizeReductionRequiredException(Exception): + pass 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 da21e1074c..fb6a833091 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 @@ -152,3 +152,51 @@ def test_interpolated_page_size_raises_on_non_integer(): config={"page_size": "invalid"}, parameters={}, ) + + +@pytest.mark.parametrize( + "initial_page_size,expected_after_reduce", + [ + pytest.param(100, 50, id="halve_100"), + pytest.param(10, 5, id="halve_10"), + pytest.param(3, 1, id="halve_3_floors_to_1"), + pytest.param(2, 1, id="halve_2_to_1"), + pytest.param(1, 1, id="already_at_minimum"), + ], +) +def test_reduce_page_size(initial_page_size, expected_after_reduce): + strategy = CursorPaginationStrategy( + page_size=initial_page_size, cursor_value="token", config={}, parameters={} + ) + strategy.reduce_page_size() + assert strategy.get_page_size() == expected_after_reduce + + +def test_reduce_page_size_multiple_times(): + strategy = CursorPaginationStrategy( + page_size=100, cursor_value="token", config={}, parameters={} + ) + strategy.reduce_page_size() + assert strategy.get_page_size() == 50 + strategy.reduce_page_size() + assert strategy.get_page_size() == 25 + strategy.reduce_page_size() + assert strategy.get_page_size() == 12 + + +def test_reset_page_size_restores_default(): + strategy = CursorPaginationStrategy( + page_size=100, cursor_value="token", config={}, parameters={} + ) + strategy.reduce_page_size() + assert strategy.get_page_size() == 50 + strategy.reset_page_size() + assert strategy.get_page_size() == 100 + + +def test_reduce_page_size_noop_when_none(): + strategy = CursorPaginationStrategy( + page_size=None, cursor_value="token", config={}, parameters={} + ) + strategy.reduce_page_size() + assert strategy.get_page_size() 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 28f6717f54..854cddaf35 100644 --- a/unit_tests/sources/declarative/requesters/paginators/test_offset_increment.py +++ b/unit_tests/sources/declarative/requesters/paginators/test_offset_increment.py @@ -146,3 +146,42 @@ def test_offset_increment_paginator_strategy_initial_token( ) assert paginator_strategy.initial_token == expected_initial_token + + +@pytest.mark.parametrize( + "initial_page_size,expected_after_reduce", + [ + pytest.param(100, 50, id="halve_100"), + pytest.param(10, 5, id="halve_10"), + pytest.param(3, 1, id="halve_3_floors_to_1"), + pytest.param(1, 1, id="already_at_minimum"), + ], +) +def test_reduce_page_size(initial_page_size, expected_after_reduce): + strategy = OffsetIncrement( + page_size=initial_page_size, parameters={}, config={}, extractor=None + ) + strategy.reduce_page_size() + assert strategy.get_page_size() == expected_after_reduce + + +def test_reduce_page_size_multiple_times(): + strategy = OffsetIncrement(page_size=100, parameters={}, config={}, extractor=None) + strategy.reduce_page_size() + assert strategy.get_page_size() == 50 + strategy.reduce_page_size() + assert strategy.get_page_size() == 25 + + +def test_reset_page_size_restores_default(): + strategy = OffsetIncrement(page_size=100, parameters={}, config={}, extractor=None) + strategy.reduce_page_size() + assert strategy.get_page_size() == 50 + strategy.reset_page_size() + assert strategy.get_page_size() == 100 + + +def test_reduce_page_size_noop_when_none(): + strategy = OffsetIncrement(page_size=None, parameters={}, config={}, extractor=None) + strategy.reduce_page_size() + assert strategy.get_page_size() is None 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 32af20b50d..0c792e09aa 100644 --- a/unit_tests/sources/declarative/requesters/paginators/test_page_increment.py +++ b/unit_tests/sources/declarative/requesters/paginators/test_page_increment.py @@ -105,3 +105,40 @@ def test_page_increment_paginator_strategy_initial_token( ) assert paginator_strategy.initial_token == expected_initial_token + + +@pytest.mark.parametrize( + "initial_page_size,expected_after_reduce", + [ + pytest.param(100, 50, id="halve_100"), + pytest.param(10, 5, id="halve_10"), + pytest.param(3, 1, id="halve_3_floors_to_1"), + pytest.param(1, 1, id="already_at_minimum"), + ], +) +def test_reduce_page_size(initial_page_size, expected_after_reduce): + strategy = PageIncrement(page_size=initial_page_size, parameters={}, config={}) + strategy.reduce_page_size() + assert strategy.get_page_size() == expected_after_reduce + + +def test_reduce_page_size_multiple_times(): + strategy = PageIncrement(page_size=100, parameters={}, config={}) + strategy.reduce_page_size() + assert strategy.get_page_size() == 50 + strategy.reduce_page_size() + assert strategy.get_page_size() == 25 + + +def test_reset_page_size_restores_default(): + strategy = PageIncrement(page_size=100, parameters={}, config={}) + strategy.reduce_page_size() + assert strategy.get_page_size() == 50 + strategy.reset_page_size() + assert strategy.get_page_size() == 100 + + +def test_reduce_page_size_noop_when_none(): + strategy = PageIncrement(page_size=None, parameters={}, config={}) + strategy.reduce_page_size() + assert strategy.get_page_size() is None diff --git a/unit_tests/sources/declarative/retrievers/test_simple_retriever.py b/unit_tests/sources/declarative/retrievers/test_simple_retriever.py index d9585dbd44..0b3949d948 100644 --- a/unit_tests/sources/declarative/retrievers/test_simple_retriever.py +++ b/unit_tests/sources/declarative/retrievers/test_simple_retriever.py @@ -34,10 +34,16 @@ GroupByKey, PropertyLimitType, ) -from airbyte_cdk.sources.declarative.requesters.request_option import RequestOptionType +from airbyte_cdk.sources.declarative.requesters.request_option import ( + RequestOption, + RequestOptionType, +) from airbyte_cdk.sources.declarative.requesters.requester import HttpMethod, Requester from airbyte_cdk.sources.declarative.retrievers.pagination_tracker import PaginationTracker from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever +from airbyte_cdk.sources.streams.http.page_size_reduction_exception import ( + PageSizeReductionRequiredException, +) from airbyte_cdk.sources.streams.http.pagination_reset_exception import ( PaginationResetRequiredException, ) @@ -1423,6 +1429,156 @@ def test_given_reach_pagination_limit_after_two_pages_when_read_records_than_red } +def test_given_page_size_reduction_exception_when_read_records_then_retry_same_page_with_reduced_size(): + """Verify that `PageSizeReductionRequiredException` causes the retriever to + reduce the page size on the pagination strategy and retry the same page + (same `next_page_token`) with the halved page size in the request, then + reset the page size after a successful fetch. + """ + strategy = CursorPaginationStrategy( + page_size=100, cursor_value="{{ response.next }}", config={}, parameters={} + ) + paginator = DefaultPaginator( + pagination_strategy=strategy, + url_base="https://api.example.com", + config={}, + parameters={}, + page_size_option=RequestOption( + inject_into=RequestOptionType.request_parameter, + field_name="page_size", + parameters={}, + ), + ) + + requester = Mock(spec=Requester) + # Page 1 succeeds → page 2 hits 502 (exception) → page 2 retried succeeds → no more pages + requester.send_request.side_effect = [ + MagicMock(), # page 1 response + PageSizeReductionRequiredException(), # page 2 fails with 502 + MagicMock(), # page 2 retry with reduced page size + ] + + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.side_effect = [ + [Record(data={"id": 1}, stream_name="test", associated_slice=A_STREAM_SLICE)], + # no select_records call for the failed request + [Record(data={"id": 2}, stream_name="test", associated_slice=A_STREAM_SLICE)], + ] + + call_count = 0 + + def patched_next_page_token(response, last_page_size, last_record, last_page_token_value): + nonlocal call_count + call_count += 1 + if call_count == 1: + return {"next_page_token": "cursor_abc"} + return None # end pagination after page 2 + + paginator.next_page_token = patched_next_page_token + + retriever = SimpleRetriever( + name="test_stream", + primary_key="id", + requester=requester, + record_selector=record_selector, + paginator=paginator, + parameters={}, + config={}, + ) + + records = list(retriever.read_records(A_RECORD_SCHEMA, A_STREAM_SLICE)) + + assert len(records) == 2 + assert requester.send_request.call_count == 3 + + # The first request used the original page token (None / initial) and page_size=100 + first_call = requester.send_request.call_args_list[0].kwargs + assert first_call.get("next_page_token") is None + assert first_call["request_params"]["page_size"] == 100 + + # The second request (which raised the exception) used cursor_abc and page_size=100 + second_call = requester.send_request.call_args_list[1].kwargs + assert second_call["next_page_token"] == {"next_page_token": "cursor_abc"} + assert second_call["request_params"]["page_size"] == 100 + + # The third request (retry) used the SAME token cursor_abc but page_size=50 + third_call = requester.send_request.call_args_list[2].kwargs + assert third_call["next_page_token"] == {"next_page_token": "cursor_abc"} + assert third_call["request_params"]["page_size"] == 50 + + # After successful retry, page size should be reset to 100 + assert strategy.get_page_size() == 100 + + +def test_page_size_is_halved_in_request_during_reduction(): + """Verify that the page size in the actual request params is halved when + the reduction exception fires, and restored after a successful fetch. + """ + strategy = CursorPaginationStrategy( + page_size=100, cursor_value="{{ response.next }}", config={}, parameters={} + ) + paginator = DefaultPaginator( + pagination_strategy=strategy, + url_base="https://api.example.com", + config={}, + parameters={}, + page_size_option=RequestOption( + inject_into=RequestOptionType.request_parameter, + field_name="page_size", + parameters={}, + ), + ) + + requester = Mock(spec=Requester) + observed_request_page_sizes = [] + + def capture_page_size_on_send(**kwargs): + observed_request_page_sizes.append(kwargs.get("request_params", {}).get("page_size")) + if len(observed_request_page_sizes) == 1: + return MagicMock() + elif len(observed_request_page_sizes) == 2: + raise PageSizeReductionRequiredException() + else: + return MagicMock() + + requester.send_request.side_effect = capture_page_size_on_send + + record_selector = Mock(spec=HttpSelector) + record_selector.select_records.side_effect = [ + [Record(data={"id": 1}, stream_name="test", associated_slice=A_STREAM_SLICE)], + [Record(data={"id": 2}, stream_name="test", associated_slice=A_STREAM_SLICE)], + ] + + call_count = 0 + + def patched_next_page_token(response, last_page_size, last_record, last_page_token_value): + nonlocal call_count + call_count += 1 + if call_count == 1: + return {"next_page_token": "cursor_abc"} + return None + + paginator.next_page_token = patched_next_page_token + + retriever = SimpleRetriever( + name="test_stream", + primary_key="id", + requester=requester, + record_selector=record_selector, + paginator=paginator, + parameters={}, + config={}, + ) + + list(retriever.read_records(A_RECORD_SCHEMA, A_STREAM_SLICE)) + + # Call 1: request has page_size=100 (page 1, succeeds) + # Call 2: request has page_size=100 (page 2, raises exception → strategy reduces to 50) + # Call 3: request has page_size=50 (page 2 retry with rebuilt request, succeeds → reset to 100) + assert observed_request_page_sizes == [100, 100, 50] + assert strategy.get_page_size() == 100 + + def _mock_paginator(): paginator = Mock(spec=Paginator) paginator.get_request_params.__name__ = "get_request_params"