diff --git a/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py b/airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py index d1f2ca41e3..e527a3cfc5 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 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 + ) + + 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..92c78cc6cc 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4368,8 +4368,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..f70c960774 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,95 @@ 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) -> 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"), + ) + + +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"}, + ) + + +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")) + + 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..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,12 +8309,10 @@ 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(): +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. """ content = _REQUEST_WINDOW_REDUCTION_STREAM.format( incremental_sync=_DATETIME_BASED_CURSOR_WITH_GRANULARITY, @@ -8331,13 +8330,15 @@ def test_given_partition_router_and_request_window_splitting_then_raise(): "", resolver.preprocess_manifest(YamlDeclarativeSource._parse(content)), {} ) - with pytest.raises(ValueError) as exception: - 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 "request_window_splitting" in str(exception.value) - assert "partition router" in str(exception.value) + 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 1a65aa0b94..292064b791 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -5332,6 +5332,191 @@ 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(): + """ + 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] + del stream["incremental_sync"]["start_time_option"] + del stream["incremental_sync"]["end_time_option"] + stream["retriever"]["partition_router"] = [dict(_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) + + +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", + [ + 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. + """ + with HttpMocker() as http_mocker: + requests = _mock_only_account_b_rejected(http_mocker) + + messages = list(_read_request_window_splitting_source(manifest)) + + 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 + ) == ["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_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]["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_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_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_half_request, 1) + http_mocker.assert_number_of_calls(account_b_second_half_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 + 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(): input_config = {} manifest = {