From 6c24b2c787842d2439024c56fce88d3d258f6998 Mon Sep 17 00:00:00 2001 From: kaizerch Date: Thu, 1 Oct 2026 10:52:07 -0700 Subject: [PATCH 1/3] feat(low-code): support request window splitting on partitioned streams ConcurrentPerPartitionCursor gains split_request_window, delegating to one state-free ConcurrentCursor built from the stream's cursor factory on first use. The factory registers RequestWindowSplitting so a CustomRetriever can declare it, and create_default_stream passes the cursor's splitter to every retriever. A CustomRetriever's request_window_splitting is not validated at load time, consistent with other custom components. --- .../concurrent_partition_cursor.py | 29 ++++ .../parsers/model_to_component_factory.py | 31 +++- .../test_concurrent_perpartitioncursor.py | 103 ++++++++++++- .../test_model_to_component_factory.py | 96 +++++++++--- .../test_concurrent_declarative_source.py | 137 ++++++++++++++++++ 5 files changed, 371 insertions(+), 25 deletions(-) diff --git a/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py b/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py index d1f2ca41e3..37c013cabf 100644 --- a/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py +++ b/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py @@ -175,6 +175,7 @@ def __init__( self._partition_key_to_index: dict[str, int] = {} self._lock = threading.Lock() + self._request_window_splitting_cursor: Optional[ConcurrentCursor] = None self._lookback_window: int = 0 self._new_global_cursor: Optional[StreamState] = None self._number_of_partitions: int = 0 @@ -677,3 +678,31 @@ def get_cursor_datetime_from_state( # Create a cursor to delegate the parsing cursor = self._cursor_factory.create(stream_state={}, runtime_lookback_window=None) return cursor.get_cursor_datetime_from_state(global_state) + + def split_request_window( + self, stream_slice: StreamSlice, min_split_window: Optional[timedelta] = None + ) -> Optional[List[StreamSlice]]: + """ + Split `stream_slice` in half the way any of this stream's per-partition cursors would. + + `ConcurrentCursor.split_request_window` only reads the slice's boundaries and the cursor's granularity + and format, never the cursor's state, and every per-partition cursor is built by the same + `_cursor_factory`. So one state-free cursor can split the slices of every partition, including after + the switch to a global cursor, when the partition's own cursor may already have been evicted. The + children keep `stream_slice`'s `partition` and `extra_fields`, and `SimpleRetriever` re-associates the + records read from them with the original slice, so `observe` and `close_partition` are unaffected. + + As for an unpartitioned stream, the children are not clamped again to the cursor's start or end. + """ + return self._get_request_window_splitting_cursor().split_request_window( + stream_slice, min_split_window + ) + + def _get_request_window_splitting_cursor(self) -> ConcurrentCursor: + # Built on first use so that streams that never split a window do not pay for an extra cursor. + with self._lock: + if self._request_window_splitting_cursor is None: + self._request_window_splitting_cursor = self._cursor_factory.create( + stream_state={}, runtime_lookback_window=None + ) + return self._request_window_splitting_cursor 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 6c8b630676..e4efd064c0 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -451,6 +451,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( RequestPath as RequestPathModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + RequestWindowSplitting as RequestWindowSplittingModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ResponseToFileExtractor as ResponseToFileExtractorModel, ) @@ -877,6 +880,7 @@ def _init_mappings(self) -> None: RemoveFieldsModel: self.create_remove_fields, RequestPathModel: self.create_request_path, RequestOptionModel: self.create_request_option, + RequestWindowSplittingModel: self.create_request_window_splitting, LegacySessionTokenAuthenticatorModel: self.create_legacy_session_token_authenticator, SelectiveAuthenticatorModel: self.create_selective_authenticator, SimpleRetrieverModel: self.create_simple_retriever, @@ -2224,6 +2228,10 @@ def create_default_stream( has_stop_condition_cursor=self._is_stop_condition_on_cursor(model), is_client_side_incremental_sync=self._is_client_side_filtering_enabled(model), cursor=concurrent_cursor, + # `create_simple_retriever` binds the cursor's splitter itself, after validating the stream. A + # `CustomRetriever` is built generically by `create_custom_component`, so this is how a + # `SimpleRetriever` subclass that declares `request_window_splitting` gets the splitter. + request_window_splitter=getattr(concurrent_cursor, "split_request_window", None), transformations=transformations, file_uploader=file_uploader, incremental_sync=model.incremental_sync, @@ -3962,7 +3970,7 @@ def _get_url(req: Requester) -> str: raise ValueError("PaginationResetLimits are not supported while having record filter.") request_window_splitting = self._create_request_window_splitting( - model, name, cursor, incremental_sync, query_properties, file_uploader + model, config, name, cursor, incremental_sync, query_properties, file_uploader ) return SimpleRetriever( name=name, @@ -4296,6 +4304,7 @@ def _uses_reduce_page_size_action(self, error_handler: Any) -> bool: def _create_request_window_splitting( self, model: SimpleRetrieverModel, + config: Config, name: str, cursor: Optional[Cursor], incremental_sync: Optional[ @@ -4338,10 +4347,21 @@ def _create_request_window_splitting( file_uploader, ) + return self.create_request_window_splitting(model.request_window_splitting, config) + + @staticmethod + def create_request_window_splitting( + model: RequestWindowSplittingModel, config: Config, **kwargs: Any + ) -> RequestWindowSplitting: + """ + Also reached directly, without `_validate_request_window_splitting_is_supported`, when a + `CustomRetriever` declares `request_window_splitting`: `create_custom_component` builds the + retriever's fields generically, without the stream's cursor or incremental sync to validate against. + """ return RequestWindowSplitting( - failure_message=model.request_window_splitting.failure_message, - min_split_window=parse_duration(model.request_window_splitting.min_split_window) - if model.request_window_splitting.min_split_window + failure_message=model.failure_message, + min_split_window=parse_duration(model.min_split_window) + if model.min_split_window else None, ) @@ -4368,8 +4388,7 @@ def _validate_request_window_splitting_is_supported( raise ValueError( f"`request_window_splitting` requires an incremental cursor that supports window splitting on " f"stream {name}. Found {type(cursor).__name__ if cursor is not None else 'no cursor'}: this is " - f"only supported today for a `DatetimeBasedCursor` with `cursor_granularity` set and no " - f"partition router combining multiple cursors." + f"only supported today for a `DatetimeBasedCursor` with `cursor_granularity` set." ) if not ( diff --git a/unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py b/unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py index 28b9b84600..8c103f302b 100644 --- a/unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py +++ b/unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py @@ -1,7 +1,7 @@ # Copyright (c) 2024 Airbyte, Inc., all rights reserved. import copy from copy import deepcopy -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Any, List, Mapping, MutableMapping, Optional, Union from unittest.mock import MagicMock, patch from urllib.parse import unquote @@ -20,13 +20,16 @@ ConcurrentDeclarativeSource, ) from airbyte_cdk.sources.declarative.incremental import ConcurrentPerPartitionCursor +from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import ( + ConcurrentCursorFactory, +) from airbyte_cdk.sources.declarative.partition_routers import ListPartitionRouter from airbyte_cdk.sources.declarative.schema import InlineSchemaLoader from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import ( DeclarativePartition, RecordCounter, ) -from airbyte_cdk.sources.streams.concurrent.cursor import CursorField +from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, CursorField from airbyte_cdk.sources.streams.concurrent.state_converters.datetime_stream_state_converter import ( CustomFormatConcurrentStreamStateConverter, ) @@ -4432,3 +4435,99 @@ def test_given_record_with_bad_cursor_value_the_global_state_parsing_does_not_br associated_slice=StreamSlice(partition={"partition_id": "1"}, cursor_slice={}), ) ) + + +def _daily_cursor_factory() -> MagicMock: + def _create_daily_cursor( + stream_state: Mapping[str, Any], runtime_lookback_window: Optional[timedelta] + ) -> ConcurrentCursor: + return ConcurrentCursor( + stream_name="test_stream", + stream_namespace=None, + stream_state=stream_state, + message_repository=MagicMock(), + connector_state_manager=MagicMock(), + connector_state_converter=CustomFormatConcurrentStreamStateConverter( + datetime_format="%Y-%m-%d", is_sequential_state=True + ), + cursor_field=CursorField(cursor_field_key="updated_at"), + slice_boundary_fields=("start_time", "end_time"), + start=datetime(2024, 1, 1, tzinfo=timezone.utc), + end_provider=lambda: datetime(2024, 2, 1, tzinfo=timezone.utc), + cursor_granularity=timedelta(days=1), + ) + + return MagicMock(wraps=ConcurrentCursorFactory(_create_daily_cursor)) + + +def _per_partition_cursor( + cursor_factory: MagicMock, use_global_cursor: bool = False +) -> ConcurrentPerPartitionCursor: + return ConcurrentPerPartitionCursor( + cursor_factory=cursor_factory, + partition_router=ListPartitionRouter( + values=["1", "2"], cursor_field="partition_id", config={}, parameters={} + ), + stream_name="test_stream", + stream_namespace=None, + stream_state={}, + message_repository=MagicMock(), + connector_state_manager=MagicMock(), + connector_state_converter=CustomFormatConcurrentStreamStateConverter( + datetime_format="%Y-%m-%d", is_sequential_state=True + ), + cursor_field=CursorField(cursor_field_key="updated_at"), + use_global_cursor=use_global_cursor, + ) + + +def _partition_window(start: str, end: str) -> StreamSlice: + return StreamSlice( + partition={"partition_id": "2"}, + cursor_slice={"start_time": start, "end_time": end}, + extra_fields={"partition_name": "second"}, + ) + + +@pytest.mark.parametrize("use_global_cursor", [False, True]) +def test_given_partition_window_when_split_request_window_then_halves_keep_partition_and_extra_fields( + use_global_cursor: bool, +) -> None: + cursor = _per_partition_cursor(_daily_cursor_factory(), use_global_cursor) + + children = cursor.split_request_window(_partition_window("2024-01-01", "2024-01-14")) + + assert children == [ + _partition_window("2024-01-01", "2024-01-07"), + _partition_window("2024-01-08", "2024-01-14"), + ] + assert [child.extra_fields for child in children] == [{"partition_name": "second"}] * 2 + + +def test_given_one_day_partition_window_when_split_request_window_then_none() -> None: + cursor = _per_partition_cursor(_daily_cursor_factory()) + + assert cursor.split_request_window(_partition_window("2024-01-01", "2024-01-01")) is None + + +def test_given_min_split_window_when_split_request_window_then_stop_at_min_split_window() -> None: + cursor = _per_partition_cursor(_daily_cursor_factory()) + + assert ( + cursor.split_request_window( + _partition_window("2024-01-01", "2024-01-03"), min_split_window=timedelta(days=2) + ) + is None + ) + + +def test_split_request_window_builds_one_state_free_cursor_on_first_use() -> None: + cursor_factory = _daily_cursor_factory() + cursor = _per_partition_cursor(cursor_factory) + assert cursor_factory.create.call_count == 0 + + cursor.split_request_window(_partition_window("2024-01-01", "2024-01-14")) + cursor.split_request_window(_partition_window("2024-01-01", "2024-01-07")) + + cursor_factory.create.assert_called_once_with(stream_state={}, runtime_lookback_window=None) + assert cursor._cursor_per_partition == {} 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 40b11eb0d9..1e87c92903 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 @@ -8308,36 +8308,98 @@ def test_given_incrementing_count_cursor_and_request_window_splitting_then_raise assert "cursor_granularity" in str(exception.value) -def test_given_partition_router_and_request_window_splitting_then_raise(): +_LIST_PARTITION_ROUTER = ( + " partition_router:\n" + " type: ListPartitionRouter\n" + ' values: ["a", "b"]\n' + " cursor_field: partition_field\n" +) + + +def _partitioned_request_window_splitting_stream(retriever_type="type: SimpleRetriever"): + content = ( + _REQUEST_WINDOW_REDUCTION_STREAM.format( + incremental_sync=_DATETIME_BASED_CURSOR_WITH_GRANULARITY, + request_window_splitting=( + "request_window_splitting:\n" + " type: RequestWindowSplitting\n" + " failure_message: Lower time_window so that each request covers less data.\n" + " min_split_window: P1D" + ), + action="SPLIT_REQUEST_WINDOW", + ) + .replace(" requester:\n", _LIST_PARTITION_ROUTER + " requester:\n") + .replace(" type: SimpleRetriever", f" {retriever_type}") + ) + stream_manifest = transformer.propagate_types_and_parameters( + "", resolver.preprocess_manifest(YamlDeclarativeSource._parse(content)), {} + ) + return factory.create_component( + model_type=DeclarativeStreamModel, component_definition=stream_manifest, config={} + ) + + +def test_given_partition_router_and_request_window_splitting_then_split_with_per_partition_cursor(): """ - A partition router combines multiple cursors (one per partition) behind a `ConcurrentPerPartitionCursor`, - which does not implement `split_request_window` itself - there is no longer one cursor whose own - granularity/boundary fields splitting could rely on. Rejected the same way a stream with no incremental - cursor at all is. + A partition router puts one cursor per partition behind a `ConcurrentPerPartitionCursor`, which splits a + window the way any of those cursors would, so the stream is accepted and the retriever is bound to it. """ + stream = _partitioned_request_window_splitting_stream() + retriever = get_retriever(stream) + + assert isinstance(stream.cursor, ConcurrentPerPartitionCursor) + assert retriever.request_window_splitting == RequestWindowSplitting( + failure_message="Lower time_window so that each request covers less data.", + min_split_window=timedelta(days=1), + ) + assert retriever.request_window_splitter.__self__ is stream.cursor + assert retriever.request_window_splitter.__name__ == "split_request_window" + + +def test_given_custom_retriever_with_request_window_splitting_then_create_retriever_with_splitter(): + """ + `create_custom_component` builds a `CustomRetriever`'s fields generically, so declaring the block on it + needs `RequestWindowSplitting` to be a known component type, and the splitter has to be handed to it by + the stream rather than bound by `create_simple_retriever`. + """ + stream = _partitioned_request_window_splitting_stream( + retriever_type=( + "type: CustomRetriever\n" + " class_name: unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever" + ) + ) + retriever = get_retriever(stream) + + assert isinstance(retriever, TestingCustomRetriever) + assert retriever.request_window_splitting == RequestWindowSplitting( + failure_message="Lower time_window so that each request covers less data.", + min_split_window=timedelta(days=1), + ) + assert retriever.request_window_splitter.__self__ is stream.cursor + + +def test_given_custom_retriever_without_incremental_sync_then_no_request_window_splitter(): content = _REQUEST_WINDOW_REDUCTION_STREAM.format( - incremental_sync=_DATETIME_BASED_CURSOR_WITH_GRANULARITY, - request_window_splitting="request_window_splitting:\n type: RequestWindowSplitting", - action="SPLIT_REQUEST_WINDOW", + incremental_sync="", + request_window_splitting="", + action="RETRY", ).replace( - " requester:\n", - " partition_router:\n" - " type: ListPartitionRouter\n" - ' values: ["a", "b"]\n' - " cursor_field: partition_field\n" - " requester:\n", + " type: SimpleRetriever", + " type: CustomRetriever\n" + " class_name: unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever", ) stream_manifest = transformer.propagate_types_and_parameters( "", resolver.preprocess_manifest(YamlDeclarativeSource._parse(content)), {} ) - with pytest.raises(ValueError) as exception: + retriever = get_retriever( factory.create_component( model_type=DeclarativeStreamModel, component_definition=stream_manifest, config={} ) + ) - assert "request_window_splitting" in str(exception.value) - assert "partition router" in str(exception.value) + assert retriever.request_window_splitting is None + assert retriever.request_window_splitter is None def test_given_query_properties_and_request_window_splitting_then_raise(): diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index 1a65aa0b94..5c984bba87 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -5332,6 +5332,143 @@ def test_given_split_request_window_action_when_read_then_final_state_reflects_f ) +_ACCOUNT_PARTITION_ROUTER = { + "type": "ListPartitionRouter", + "values": ["a", "b"], + "cursor_field": "account", +} + + +def _partitioned_request_window_splitting_manifest(): + """Two accounts, each read as one 14-day window and checkpointed with its own cursor.""" + manifest = _request_window_splitting_manifest() + stream = manifest["streams"][0] + stream["incremental_sync"].update( + { + "start_datetime": "2024-01-01", + "end_datetime": "2024-01-14", + "step": "P14D", + "cursor_granularity": "P1D", + "datetime_format": "%Y-%m-%d", + } + ) + stream["retriever"]["partition_router"] = { + **_ACCOUNT_PARTITION_ROUTER, + "request_option": { + "type": "RequestOption", + "inject_into": "request_parameter", + "field_name": "account", + }, + } + return manifest + + +def _partitioned_request_window_splitting_manifest_with_interpolated_window(): + """ + The shape of TikTok Marketing's daily report streams: a list of partition routers, which the factory wraps + in a `CartesianProductStreamSlicer`, and request parameters that read the window from `stream_interval` + rather than injecting it through `start_time_option`/`end_time_option`. + """ + manifest = _partitioned_request_window_splitting_manifest() + stream = manifest["streams"][0] + del stream["incremental_sync"]["start_time_option"] + del stream["incremental_sync"]["end_time_option"] + stream["retriever"]["partition_router"] = [_ACCOUNT_PARTITION_ROUTER] + stream["retriever"]["requester"]["request_parameters"] = { + "account": "{{ stream_partition.account }}", + "start": "{{ stream_interval.start_time }}", + "end": "{{ stream_interval.end_time }}", + } + return manifest + + +def _account_window_request(account: str, start: str, end: str) -> HttpRequest: + return HttpRequest( + "https://example.org/test", + query_params={"account": account, "start": start, "end": end}, + ) + + +def _items_response(*updated_at: str) -> HttpResponse: + return HttpResponse(json.dumps({"items": [{"updated_at": value} for value in updated_at]}), 200) + + +@pytest.mark.parametrize( + "manifest", + [ + pytest.param(_partitioned_request_window_splitting_manifest(), id="request_options"), + pytest.param( + _partitioned_request_window_splitting_manifest_with_interpolated_window(), + id="interpolated_window", + ), + ], +) +def test_given_partitioned_stream_when_one_partition_window_is_rejected_then_split_only_that_partition( + manifest, +): + """ + The Google Ads and TikTok case: only one account's window is too large. That account's window is read as + two halves, the other account is read once, and each account is checkpointed with its own cursor. + """ + account_a_request = _account_window_request("a", "2024-01-01", "2024-01-14") + account_b_request = _account_window_request("b", "2024-01-01", "2024-01-14") + account_b_first_half_request = _account_window_request("b", "2024-01-01", "2024-01-07") + account_b_second_half_request = _account_window_request("b", "2024-01-08", "2024-01-14") + + with HttpMocker() as http_mocker: + http_mocker.get(account_a_request, _items_response("2024-01-10")) + http_mocker.get(account_b_request, HttpResponse("", 400)) + http_mocker.get(account_b_first_half_request, _items_response("2024-01-03")) + http_mocker.get(account_b_second_half_request, _items_response("2024-01-12")) + + messages = list(_read_request_window_splitting_source(manifest)) + + http_mocker.assert_number_of_calls(account_a_request, 1) + http_mocker.assert_number_of_calls(account_b_request, 1) + http_mocker.assert_number_of_calls(account_b_first_half_request, 1) + http_mocker.assert_number_of_calls(account_b_second_half_request, 1) + + assert sorted( + message.record.data["updated_at"] for message in messages if message.type == Type.RECORD + ) == ["2024-01-03", "2024-01-10", "2024-01-12"] + final_state = get_states_for_stream(stream_name="Test", messages=messages)[ + -1 + ].stream.stream_state.__dict__ + assert final_state["use_global_cursor"] is False + assert sorted(final_state["states"], key=lambda state: state["partition"]["account"]) == [ + {"partition": {"account": "a"}, "cursor": {"updated_at": "2024-01-10"}}, + {"partition": {"account": "b"}, "cursor": {"updated_at": "2024-01-12"}}, + ] + + +def test_given_partitioned_stream_when_one_day_window_is_rejected_then_transient_error(): + account_a_request = _account_window_request("a", "2024-01-01", "2024-01-02") + account_b_request = _account_window_request("b", "2024-01-01", "2024-01-02") + account_b_first_day_request = _account_window_request("b", "2024-01-01", "2024-01-01") + manifest = _partitioned_request_window_splitting_manifest() + manifest["streams"][0]["incremental_sync"]["end_datetime"] = "2024-01-02" + + messages = [] + with HttpMocker() as http_mocker: + http_mocker.get(account_a_request, _items_response("2024-01-01")) + http_mocker.get(account_b_request, HttpResponse("", 400)) + http_mocker.get(account_b_first_day_request, HttpResponse("", 400)) + + with pytest.raises(AirbyteTracedException): + messages.extend(_read_request_window_splitting_source(manifest)) + + http_mocker.assert_number_of_calls(account_b_request, 1) + http_mocker.assert_number_of_calls(account_b_first_day_request, 1) + + errors = [ + message.trace.error + for message in messages + if message.type == Type.TRACE and message.trace.type == TraceType.ERROR + ] + assert errors[0].failure_type == FailureType.transient_error + assert "could not split its request window" in errors[0].internal_message + + def test_given_pagination_limit_reached_when_read_then_reset_pagination(): input_config = {} manifest = { From e0371aac9c3844ed9c3c2fa84eab16ede26b6e2f Mon Sep 17 00:00:00 2001 From: kaizerch Date: Thu, 1 Oct 2026 11:39:11 -0700 Subject: [PATCH 2/3] fixes based on cdk-pr-review --- .../concurrent_partition_cursor.py | 2 +- .../parsers/model_to_component_factory.py | 28 +----- .../test_concurrent_perpartitioncursor.py | 14 +-- .../test_model_to_component_factory.py | 95 ++++--------------- .../test_concurrent_declarative_source.py | 92 +++++++++++++----- 5 files changed, 97 insertions(+), 134 deletions(-) diff --git a/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py b/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py index 37c013cabf..e527a3cfc5 100644 --- a/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py +++ b/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py @@ -692,7 +692,7 @@ def split_request_window( children keep `stream_slice`'s `partition` and `extra_fields`, and `SimpleRetriever` re-associates the records read from them with the original slice, so `observe` and `close_partition` are unaffected. - As for an unpartitioned stream, the children are not clamped again to the cursor's start or end. + As with an unpartitioned stream, the children are not clamped again to the cursor's start or end. """ return self._get_request_window_splitting_cursor().split_request_window( stream_slice, min_split_window 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 e4efd064c0..92c78cc6cc 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -451,9 +451,6 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( RequestPath as RequestPathModel, ) -from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( - RequestWindowSplitting as RequestWindowSplittingModel, -) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ResponseToFileExtractor as ResponseToFileExtractorModel, ) @@ -880,7 +877,6 @@ def _init_mappings(self) -> None: RemoveFieldsModel: self.create_remove_fields, RequestPathModel: self.create_request_path, RequestOptionModel: self.create_request_option, - RequestWindowSplittingModel: self.create_request_window_splitting, LegacySessionTokenAuthenticatorModel: self.create_legacy_session_token_authenticator, SelectiveAuthenticatorModel: self.create_selective_authenticator, SimpleRetrieverModel: self.create_simple_retriever, @@ -2228,10 +2224,6 @@ def create_default_stream( has_stop_condition_cursor=self._is_stop_condition_on_cursor(model), is_client_side_incremental_sync=self._is_client_side_filtering_enabled(model), cursor=concurrent_cursor, - # `create_simple_retriever` binds the cursor's splitter itself, after validating the stream. A - # `CustomRetriever` is built generically by `create_custom_component`, so this is how a - # `SimpleRetriever` subclass that declares `request_window_splitting` gets the splitter. - request_window_splitter=getattr(concurrent_cursor, "split_request_window", None), transformations=transformations, file_uploader=file_uploader, incremental_sync=model.incremental_sync, @@ -3970,7 +3962,7 @@ def _get_url(req: Requester) -> str: raise ValueError("PaginationResetLimits are not supported while having record filter.") request_window_splitting = self._create_request_window_splitting( - model, config, name, cursor, incremental_sync, query_properties, file_uploader + model, name, cursor, incremental_sync, query_properties, file_uploader ) return SimpleRetriever( name=name, @@ -4304,7 +4296,6 @@ def _uses_reduce_page_size_action(self, error_handler: Any) -> bool: def _create_request_window_splitting( self, model: SimpleRetrieverModel, - config: Config, name: str, cursor: Optional[Cursor], incremental_sync: Optional[ @@ -4347,21 +4338,10 @@ def _create_request_window_splitting( file_uploader, ) - return self.create_request_window_splitting(model.request_window_splitting, config) - - @staticmethod - def create_request_window_splitting( - model: RequestWindowSplittingModel, config: Config, **kwargs: Any - ) -> RequestWindowSplitting: - """ - Also reached directly, without `_validate_request_window_splitting_is_supported`, when a - `CustomRetriever` declares `request_window_splitting`: `create_custom_component` builds the - retriever's fields generically, without the stream's cursor or incremental sync to validate against. - """ return RequestWindowSplitting( - failure_message=model.failure_message, - min_split_window=parse_duration(model.min_split_window) - if model.min_split_window + failure_message=model.request_window_splitting.failure_message, + min_split_window=parse_duration(model.request_window_splitting.min_split_window) + if model.request_window_splitting.min_split_window else None, ) diff --git a/unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py b/unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py index 8c103f302b..f70c960774 100644 --- a/unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py +++ b/unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py @@ -4460,9 +4460,7 @@ def _create_daily_cursor( return MagicMock(wraps=ConcurrentCursorFactory(_create_daily_cursor)) -def _per_partition_cursor( - cursor_factory: MagicMock, use_global_cursor: bool = False -) -> ConcurrentPerPartitionCursor: +def _per_partition_cursor(cursor_factory: MagicMock) -> ConcurrentPerPartitionCursor: return ConcurrentPerPartitionCursor( cursor_factory=cursor_factory, partition_router=ListPartitionRouter( @@ -4477,7 +4475,6 @@ def _per_partition_cursor( datetime_format="%Y-%m-%d", is_sequential_state=True ), cursor_field=CursorField(cursor_field_key="updated_at"), - use_global_cursor=use_global_cursor, ) @@ -4489,11 +4486,10 @@ def _partition_window(start: str, end: str) -> StreamSlice: ) -@pytest.mark.parametrize("use_global_cursor", [False, True]) -def test_given_partition_window_when_split_request_window_then_halves_keep_partition_and_extra_fields( - use_global_cursor: bool, -) -> None: - cursor = _per_partition_cursor(_daily_cursor_factory(), use_global_cursor) +def test_given_partition_window_when_split_request_window_then_halves_keep_partition_and_extra_fields() -> ( + None +): + cursor = _per_partition_cursor(_daily_cursor_factory()) children = cursor.split_request_window(_partition_window("2024-01-01", "2024-01-14")) 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 1e87c92903..340259a084 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 @@ -8198,6 +8198,7 @@ def test_given_no_incremental_sync_and_request_window_splitting_then_raise(): _request_window_splitting_stream(incremental_sync="") assert "cursor_granularity" in str(exception.value) + assert "partition router" not in str(exception.value) def test_given_datetime_based_cursor_without_cursor_granularity_then_raise(): @@ -8308,98 +8309,36 @@ def test_given_incrementing_count_cursor_and_request_window_splitting_then_raise assert "cursor_granularity" in str(exception.value) -_LIST_PARTITION_ROUTER = ( - " partition_router:\n" - " type: ListPartitionRouter\n" - ' values: ["a", "b"]\n' - " cursor_field: partition_field\n" -) - - -def _partitioned_request_window_splitting_stream(retriever_type="type: SimpleRetriever"): - content = ( - _REQUEST_WINDOW_REDUCTION_STREAM.format( - incremental_sync=_DATETIME_BASED_CURSOR_WITH_GRANULARITY, - request_window_splitting=( - "request_window_splitting:\n" - " type: RequestWindowSplitting\n" - " failure_message: Lower time_window so that each request covers less data.\n" - " min_split_window: P1D" - ), - action="SPLIT_REQUEST_WINDOW", - ) - .replace(" requester:\n", _LIST_PARTITION_ROUTER + " requester:\n") - .replace(" type: SimpleRetriever", f" {retriever_type}") - ) - stream_manifest = transformer.propagate_types_and_parameters( - "", resolver.preprocess_manifest(YamlDeclarativeSource._parse(content)), {} - ) - return factory.create_component( - model_type=DeclarativeStreamModel, component_definition=stream_manifest, config={} - ) - - def test_given_partition_router_and_request_window_splitting_then_split_with_per_partition_cursor(): """ A partition router puts one cursor per partition behind a `ConcurrentPerPartitionCursor`, which splits a window the way any of those cursors would, so the stream is accepted and the retriever is bound to it. """ - stream = _partitioned_request_window_splitting_stream() - retriever = get_retriever(stream) - - assert isinstance(stream.cursor, ConcurrentPerPartitionCursor) - assert retriever.request_window_splitting == RequestWindowSplitting( - failure_message="Lower time_window so that each request covers less data.", - min_split_window=timedelta(days=1), - ) - assert retriever.request_window_splitter.__self__ is stream.cursor - assert retriever.request_window_splitter.__name__ == "split_request_window" - - -def test_given_custom_retriever_with_request_window_splitting_then_create_retriever_with_splitter(): - """ - `create_custom_component` builds a `CustomRetriever`'s fields generically, so declaring the block on it - needs `RequestWindowSplitting` to be a known component type, and the splitter has to be handed to it by - the stream rather than bound by `create_simple_retriever`. - """ - stream = _partitioned_request_window_splitting_stream( - retriever_type=( - "type: CustomRetriever\n" - " class_name: unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever" - ) - ) - retriever = get_retriever(stream) - - assert isinstance(retriever, TestingCustomRetriever) - assert retriever.request_window_splitting == RequestWindowSplitting( - failure_message="Lower time_window so that each request covers less data.", - min_split_window=timedelta(days=1), - ) - assert retriever.request_window_splitter.__self__ is stream.cursor - - -def test_given_custom_retriever_without_incremental_sync_then_no_request_window_splitter(): content = _REQUEST_WINDOW_REDUCTION_STREAM.format( - incremental_sync="", - request_window_splitting="", - action="RETRY", + incremental_sync=_DATETIME_BASED_CURSOR_WITH_GRANULARITY, + request_window_splitting="request_window_splitting:\n type: RequestWindowSplitting", + action="SPLIT_REQUEST_WINDOW", ).replace( - " type: SimpleRetriever", - " type: CustomRetriever\n" - " class_name: unit_tests.sources.declarative.parsers.testing_components.TestingCustomRetriever", + " requester:\n", + " partition_router:\n" + " type: ListPartitionRouter\n" + ' values: ["a", "b"]\n' + " cursor_field: partition_field\n" + " requester:\n", ) stream_manifest = transformer.propagate_types_and_parameters( "", resolver.preprocess_manifest(YamlDeclarativeSource._parse(content)), {} ) - retriever = get_retriever( - factory.create_component( - model_type=DeclarativeStreamModel, component_definition=stream_manifest, config={} - ) + stream = factory.create_component( + model_type=DeclarativeStreamModel, component_definition=stream_manifest, config={} ) + retriever = get_retriever(stream) - assert retriever.request_window_splitting is None - assert retriever.request_window_splitter is None + assert isinstance(stream.cursor, ConcurrentPerPartitionCursor) + assert retriever.request_window_splitting == RequestWindowSplitting() + assert retriever.request_window_splitter.__self__ is stream.cursor + assert retriever.request_window_splitter.__name__ == "split_request_window" def test_given_query_properties_and_request_window_splitting_then_raise(): diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index 5c984bba87..f9c87c6af0 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -5373,7 +5373,7 @@ def _partitioned_request_window_splitting_manifest_with_interpolated_window(): stream = manifest["streams"][0] del stream["incremental_sync"]["start_time_option"] del stream["incremental_sync"]["end_time_option"] - stream["retriever"]["partition_router"] = [_ACCOUNT_PARTITION_ROUTER] + stream["retriever"]["partition_router"] = [dict(_ACCOUNT_PARTITION_ROUTER)] stream["retriever"]["requester"]["request_parameters"] = { "account": "{{ stream_partition.account }}", "start": "{{ stream_interval.start_time }}", @@ -5393,6 +5393,19 @@ def _items_response(*updated_at: str) -> HttpResponse: return HttpResponse(json.dumps({"items": [{"updated_at": value} for value in updated_at]}), 200) +def _mock_only_account_b_rejected(http_mocker: HttpMocker) -> List[HttpRequest]: + """Account b's 14-day window is rejected and its two halves succeed; account a's window succeeds.""" + responses = [ + (_account_window_request("a", "2024-01-01", "2024-01-14"), _items_response("2024-01-10")), + (_account_window_request("b", "2024-01-01", "2024-01-14"), HttpResponse("", 400)), + (_account_window_request("b", "2024-01-01", "2024-01-07"), _items_response("2024-01-03")), + (_account_window_request("b", "2024-01-08", "2024-01-14"), _items_response("2024-01-12")), + ] + for request, response in responses: + http_mocker.get(request, response) + return [request for request, _ in responses] + + @pytest.mark.parametrize( "manifest", [ @@ -5410,23 +5423,13 @@ def test_given_partitioned_stream_when_one_partition_window_is_rejected_then_spl The Google Ads and TikTok case: only one account's window is too large. That account's window is read as two halves, the other account is read once, and each account is checkpointed with its own cursor. """ - account_a_request = _account_window_request("a", "2024-01-01", "2024-01-14") - account_b_request = _account_window_request("b", "2024-01-01", "2024-01-14") - account_b_first_half_request = _account_window_request("b", "2024-01-01", "2024-01-07") - account_b_second_half_request = _account_window_request("b", "2024-01-08", "2024-01-14") - with HttpMocker() as http_mocker: - http_mocker.get(account_a_request, _items_response("2024-01-10")) - http_mocker.get(account_b_request, HttpResponse("", 400)) - http_mocker.get(account_b_first_half_request, _items_response("2024-01-03")) - http_mocker.get(account_b_second_half_request, _items_response("2024-01-12")) + requests = _mock_only_account_b_rejected(http_mocker) messages = list(_read_request_window_splitting_source(manifest)) - http_mocker.assert_number_of_calls(account_a_request, 1) - http_mocker.assert_number_of_calls(account_b_request, 1) - http_mocker.assert_number_of_calls(account_b_first_half_request, 1) - http_mocker.assert_number_of_calls(account_b_second_half_request, 1) + for request in requests: + http_mocker.assert_number_of_calls(request, 1) assert sorted( message.record.data["updated_at"] for message in messages if message.type == Type.RECORD @@ -5441,24 +5444,61 @@ def test_given_partitioned_stream_when_one_partition_window_is_rejected_then_spl ] -def test_given_partitioned_stream_when_one_day_window_is_rejected_then_transient_error(): - account_a_request = _account_window_request("a", "2024-01-01", "2024-01-02") - account_b_request = _account_window_request("b", "2024-01-01", "2024-01-02") - account_b_first_day_request = _account_window_request("b", "2024-01-01", "2024-01-01") +def test_given_partitioned_stream_over_switch_to_global_limit_when_window_is_rejected_then_split(): + """ + Past `SWITCH_TO_GLOBAL_LIMIT` partitions, the stream keeps a single global cursor instead of one per + partition. Splitting does not depend on the per-partition cursors, so it works the same way. + """ + from airbyte_cdk.sources.declarative.incremental import ConcurrentPerPartitionCursor + + with HttpMocker() as http_mocker: + requests = _mock_only_account_b_rejected(http_mocker) + + with patch.object(ConcurrentPerPartitionCursor, "SWITCH_TO_GLOBAL_LIMIT", 0): + messages = list( + _read_request_window_splitting_source( + _partitioned_request_window_splitting_manifest() + ) + ) + + for request in requests: + http_mocker.assert_number_of_calls(request, 1) + + final_state = get_states_for_stream(stream_name="Test", messages=messages)[ + -1 + ].stream.stream_state.__dict__ + assert final_state["use_global_cursor"] is True + assert final_state["state"] == {"updated_at": "2024-01-12"} + + +def test_given_split_window_still_rejected_when_read_then_partition_is_not_checkpointed(): + """ + Account b's first half is read, but its second half is rejected again and is already at + `min_split_window`, so the stream fails with a `transient_error`. Account b's cursor must stay at its start + instead of moving to the record read from the first half, so the next sync re-reads the whole window, + while account a is still checkpointed. + """ manifest = _partitioned_request_window_splitting_manifest() - manifest["streams"][0]["incremental_sync"]["end_datetime"] = "2024-01-02" + manifest["streams"][0]["retriever"]["request_window_splitting"]["min_split_window"] = "P7D" + account_b_request = _account_window_request("b", "2024-01-01", "2024-01-14") + account_b_first_half_request = _account_window_request("b", "2024-01-01", "2024-01-07") + account_b_second_half_request = _account_window_request("b", "2024-01-08", "2024-01-14") messages = [] with HttpMocker() as http_mocker: - http_mocker.get(account_a_request, _items_response("2024-01-01")) + http_mocker.get( + _account_window_request("a", "2024-01-01", "2024-01-14"), _items_response("2024-01-10") + ) http_mocker.get(account_b_request, HttpResponse("", 400)) - http_mocker.get(account_b_first_day_request, HttpResponse("", 400)) + http_mocker.get(account_b_first_half_request, _items_response("2024-01-03")) + http_mocker.get(account_b_second_half_request, HttpResponse("", 400)) with pytest.raises(AirbyteTracedException): messages.extend(_read_request_window_splitting_source(manifest)) http_mocker.assert_number_of_calls(account_b_request, 1) - http_mocker.assert_number_of_calls(account_b_first_day_request, 1) + http_mocker.assert_number_of_calls(account_b_first_half_request, 1) + http_mocker.assert_number_of_calls(account_b_second_half_request, 1) errors = [ message.trace.error @@ -5467,6 +5507,14 @@ def test_given_partitioned_stream_when_one_day_window_is_rejected_then_transient ] assert errors[0].failure_type == FailureType.transient_error assert "could not split its request window" in errors[0].internal_message + final_state = get_states_for_stream(stream_name="Test", messages=messages)[ + -1 + ].stream.stream_state.__dict__ + assert sorted(final_state["states"], key=lambda state: state["partition"]["account"]) == [ + {"partition": {"account": "a"}, "cursor": {"updated_at": "2024-01-10"}}, + {"partition": {"account": "b"}, "cursor": {"updated_at": "2024-01-01"}}, + ] + assert "state" not in final_state def test_given_pagination_limit_reached_when_read_then_reset_pagination(): From 1d62b0f597c3061f58ee4c8b3adde24512ff6438 Mon Sep 17 00:00:00 2001 From: kaizerch Date: Thu, 1 Oct 2026 12:14:29 -0700 Subject: [PATCH 3/3] tests: describe the interpolated-window case as a list-form partition router --- .../declarative/test_concurrent_declarative_source.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index f9c87c6af0..292064b791 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -5365,9 +5365,9 @@ def _partitioned_request_window_splitting_manifest(): def _partitioned_request_window_splitting_manifest_with_interpolated_window(): """ - The shape of TikTok Marketing's daily report streams: a list of partition routers, which the factory wraps - in a `CartesianProductStreamSlicer`, and request parameters that read the window from `stream_interval` - rather than injecting it through `start_time_option`/`end_time_option`. + A list-form partition router, which the factory wraps in a `CartesianProductStreamSlicer`, and request + parameters that read the window from `stream_interval` rather than injecting it through + `start_time_option`/`end_time_option`. """ manifest = _partitioned_request_window_splitting_manifest() stream = manifest["streams"][0]