Force global state for GroupingPartitionRouter

Author: tolik0

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -2894,7 +2894,7 @@ definitions:
         title: Lazy Read Pointer
         description: If set, this will enable lazy reading, using the initial read of parent records to extract child records.
         type: array
-        default: [ ]
+        default: []
         items:
           - type: string
         interpolation_context:
@@ -3199,7 +3199,7 @@ definitions:
     properties:
       type:
         type: string
-        enum: [ StateDelegatingStream ]
+        enum: [StateDelegatingStream]
       name:
         title: Name
         description: The stream name.

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -79,6 +79,7 @@ def __init__(
         connector_state_manager: ConnectorStateManager,
         connector_state_converter: AbstractStreamStateConverter,
         cursor_field: CursorField,
+        use_global_cursor: bool = False,
     ) -> None:
         self._global_cursor: Optional[StreamState] = {}
         self._stream_name = stream_name
@@ -106,7 +107,7 @@ def __init__(
         self._lookback_window: int = 0
         self._parent_state: Optional[StreamState] = None
         self._number_of_partitions: int = 0
-        self._use_global_cursor: bool = False
+        self._use_global_cursor: bool = use_global_cursor
         self._partition_serializer = PerPartitionKeySerializer()
         # Track the last time a state message was emitted
         self._last_emission_time: float = 0.0

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -1360,6 +1360,9 @@ def create_concurrent_cursor_from_perpartition_cursor(
         )
         stream_state = self.apply_stream_state_migrations(stream_state_migrations, stream_state)
 
+        # Per-partition state doesn't make sense for GroupingPartitionRouter, so force the global state
+        use_global_cursor = isinstance(partition_router, GroupingPartitionRouter)
+
         # Return the concurrent cursor and state converter
         return ConcurrentPerPartitionCursor(
             cursor_factory=cursor_factory,
@@ -1371,6 +1374,7 @@ def create_concurrent_cursor_from_perpartition_cursor(
             connector_state_manager=state_manager,
             connector_state_converter=connector_state_converter,
             cursor_field=cursor_field,
+            use_global_cursor=use_global_cursor,
         )
 
     @staticmethod

unit_tests/sources/declarative/partition_routers/test_grouping_partition_router.py

@@ -312,21 +312,6 @@ def test_set_initial_state_delegation(mock_config, mock_underlying_router):
     mock_underlying_router.set_initial_state.assert_called_once_with(mock_state)
 
 
-def test_get_stream_state_delegation(mock_config, mock_underlying_router):
-    """Test that get_stream_state delegates to the underlying router."""
-    router = GroupingPartitionRouter(
-        group_size=2,
-        underlying_partition_router=mock_underlying_router,
-        config=mock_config,
-    )
-    mock_state = {"state_key": "state_value"}
-    mock_underlying_router.get_stream_state = MagicMock(return_value=mock_state)
-
-    state = router.get_stream_state()
-    assert state == mock_state
-    mock_underlying_router.get_stream_state.assert_called_once()
-
-
 def test_stream_slices_extra_fields_varied(mock_config):
     """Test grouping with varied extra fields across partitions."""
     parent_stream = MockStream(
@@ -458,3 +443,72 @@ def test_get_request_params_default(mock_config, mock_underlying_router):
         )
     )
     assert params == {}
+
+
+def test_stream_slices_resume_from_state(mock_config, mock_underlying_router):
+    """Test that stream_slices resumes correctly from a previous state."""
+
+    # Simulate underlying router state handling
+    class MockPartitionRouter:
+        def __init__(self):
+            self.slices = [
+                StreamSlice(
+                    partition={"board_ids": i},
+                    cursor_slice={},
+                    extra_fields={"name": f"Board {i}", "owner": f"User{i}"},
+                )
+                for i in range(5)
+            ]
+            self.state = {"last_board_id": 0}  # Initial state
+
+        def set_initial_state(self, state):
+            self.state = state
+
+        def get_stream_state(self):
+            return self.state
+
+        def stream_slices(self):
+            last_board_id = self.state.get("last_board_id", -1)
+            for slice in self.slices:
+                board_id = slice.partition["board_ids"]
+                if board_id <= last_board_id:
+                    continue
+                self.state = {"last_board_id": board_id}
+                yield slice
+
+    underlying_router = MockPartitionRouter()
+    router = GroupingPartitionRouter(
+        group_size=2,
+        underlying_partition_router=underlying_router,
+        config=mock_config,
+        deduplicate=True,
+    )
+
+    # First sync: process first two slices
+    router.set_initial_state({"last_board_id": 0})
+    slices_iter = router.stream_slices()
+    first_batch = next(slices_iter)
+    assert first_batch == StreamSlice(
+        partition={"board_ids": [1, 2]},
+        cursor_slice={},
+        extra_fields={"name": ["Board 1", "Board 2"], "owner": ["User1", "User2"]},
+    )
+    state_after_first = router.get_stream_state()
+    assert state_after_first == {"last_board_id": 2}, "State should reflect last processed board_id"
+
+    # Simulate a new sync resuming from the previous state
+    new_router = GroupingPartitionRouter(
+        group_size=2,
+        underlying_partition_router=MockPartitionRouter(),
+        config=mock_config,
+        deduplicate=True,
+    )
+    new_router.set_initial_state(state_after_first)
+    resumed_slices = list(new_router.stream_slices())
+    assert resumed_slices == [
+        StreamSlice(
+            partition={"board_ids": [3, 4]},
+            cursor_slice={},
+            extra_fields={"name": ["Board 3", "Board 4"], "owner": ["User3", "User4"]},
+        )
+    ], "Should resume from board_id 3"