feat(low-code): support request window splitting on partitioned streams

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

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 (

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 == {}

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():