From 366f98a06d759767106dcbdb0e7c82efce0764fb Mon Sep 17 00:00:00 2001 From: brianjlai Date: Mon, 6 Jan 2025 11:31:48 -0800 Subject: [PATCH 1/3] fix offset strategy and stop condition decorator to reset the token back to 0 properly --- .../paginators/strategies/offset_increment.py | 4 +- .../paginators/strategies/stop_condition.py | 5 ++- .../paginators/test_default_paginator.py | 39 +++++++++++++++++++ .../paginators/test_offset_increment.py | 5 +-- 4 files changed, 47 insertions(+), 6 deletions(-) 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 37ba3bbfa..7b17c81b4 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/offset_increment.py @@ -83,7 +83,9 @@ def next_page_token( return self._offset def reset(self, reset_value: Optional[Any] = 0) -> None: - if not isinstance(reset_value, int): + if reset_value is None: + self._offset = 0 + elif not isinstance(reset_value, int): raise ValueError( f"Reset value {reset_value} for OffsetIncrement pagination strategy was not an integer" ) diff --git a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/stop_condition.py b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/stop_condition.py index 7722c5e73..a3c977f18 100644 --- a/airbyte_cdk/sources/declarative/requesters/paginators/strategies/stop_condition.py +++ b/airbyte_cdk/sources/declarative/requesters/paginators/strategies/stop_condition.py @@ -53,7 +53,10 @@ def next_page_token( return self._delegate.next_page_token(response, last_page_size, last_record) def reset(self, reset_value: Optional[Any] = None) -> None: - self._delegate.reset(reset_value) + if reset_value: + self._delegate.reset(reset_value) + else: + self._delegate.reset() def get_page_size(self) -> Optional[int]: return self._delegate.get_page_size() diff --git a/unit_tests/sources/declarative/requesters/paginators/test_default_paginator.py b/unit_tests/sources/declarative/requesters/paginators/test_default_paginator.py index 1cd34c42f..f315abb9e 100644 --- a/unit_tests/sources/declarative/requesters/paginators/test_default_paginator.py +++ b/unit_tests/sources/declarative/requesters/paginators/test_default_paginator.py @@ -9,6 +9,7 @@ import requests from airbyte_cdk.sources.declarative.decoders import JsonDecoder, XmlDecoder +from airbyte_cdk.sources.declarative.incremental import DatetimeBasedCursor from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean from airbyte_cdk.sources.declarative.requesters.paginators.default_paginator import ( DefaultPaginator, @@ -22,6 +23,10 @@ from airbyte_cdk.sources.declarative.requesters.paginators.strategies.offset_increment import ( OffsetIncrement, ) +from airbyte_cdk.sources.declarative.requesters.paginators.strategies.stop_condition import ( + CursorStopCondition, + StopConditionPaginationStrategyDecorator, +) from airbyte_cdk.sources.declarative.requesters.request_path import RequestPath @@ -356,6 +361,40 @@ def test_reset(inject_on_first_request): assert request_parameters_for_second_request != request_parameters_after_reset +def test_data_feed_paginator_with_stop_page_condition(): + config = {} + + cursor = DatetimeBasedCursor( + cursor_field="updated_at", + datetime_format="%Y-%m-$d", + start_datetime="2024-01-01", + config=config, + parameters={}, + ) + + wrapped_strategy = StopConditionPaginationStrategyDecorator( + _delegate=OffsetIncrement( + config={}, page_size=2, inject_on_first_request=False, parameters={} + ), + stop_condition=CursorStopCondition(cursor=cursor), + ) + + paginator = DefaultPaginator( + pagination_strategy=wrapped_strategy, + config=config, + url_base="https://airbyte.io", + parameters={}, + page_size_option=RequestOption( + inject_into=RequestOptionType.request_parameter, field_name="limit", parameters={} + ), + page_token_option=RequestOption( + inject_into=RequestOptionType.request_parameter, field_name="offset", parameters={} + ), + ) + + paginator.reset() + + def test_initial_token_with_offset_pagination(): page_size_request_option = RequestOption( inject_into=RequestOptionType.request_parameter, field_name="limit", parameters={} 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 d443132ed..692b80272 100644 --- a/unit_tests/sources/declarative/requesters/paginators/test_offset_increment.py +++ b/unit_tests/sources/declarative/requesters/paginators/test_offset_increment.py @@ -104,8 +104,5 @@ def test_offset_increment_reset(reset_value, expected_initial_token, expected_er with pytest.raises(expected_error): paginator_strategy.reset(reset_value=reset_value) else: - if reset_value is None: - paginator_strategy.reset() - else: - paginator_strategy.reset(reset_value=reset_value) + paginator_strategy.reset(reset_value=reset_value) assert paginator_strategy.initial_token == expected_initial_token From 356e32129ccbb2b98c93c2b3c9987caa6e27ef0e Mon Sep 17 00:00:00 2001 From: brianjlai Date: Mon, 6 Jan 2025 12:01:08 -0800 Subject: [PATCH 2/3] date format in test --- .../declarative/requesters/paginators/test_default_paginator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unit_tests/sources/declarative/requesters/paginators/test_default_paginator.py b/unit_tests/sources/declarative/requesters/paginators/test_default_paginator.py index f315abb9e..fcf631f7f 100644 --- a/unit_tests/sources/declarative/requesters/paginators/test_default_paginator.py +++ b/unit_tests/sources/declarative/requesters/paginators/test_default_paginator.py @@ -366,7 +366,7 @@ def test_data_feed_paginator_with_stop_page_condition(): cursor = DatetimeBasedCursor( cursor_field="updated_at", - datetime_format="%Y-%m-$d", + datetime_format="%Y-%m-%d", start_datetime="2024-01-01", config=config, parameters={}, From 62c707f16cd39a5dccb05e0b9010e17b355485c8 Mon Sep 17 00:00:00 2001 From: brianjlai Date: Mon, 6 Jan 2025 12:13:56 -0800 Subject: [PATCH 3/3] fix test --- .../declarative/requesters/paginators/test_stop_condition.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unit_tests/sources/declarative/requesters/paginators/test_stop_condition.py b/unit_tests/sources/declarative/requesters/paginators/test_stop_condition.py index ea1d38e24..88da84589 100644 --- a/unit_tests/sources/declarative/requesters/paginators/test_stop_condition.py +++ b/unit_tests/sources/declarative/requesters/paginators/test_stop_condition.py @@ -108,7 +108,7 @@ def test_when_reset_then_delegate(mocked_pagination_strategy, mocked_stop_condit mocked_pagination_strategy, mocked_stop_condition ) decorator.reset() - mocked_pagination_strategy.reset.assert_called_once_with(None) + mocked_pagination_strategy.reset.assert_called_once_with() def test_when_get_page_size_then_delegate(mocked_pagination_strategy, mocked_stop_condition):