Add skipping duplicate partitions

Author: tolik0

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -6,10 +6,10 @@
 import logging
 import threading
 import time
-from collections import OrderedDict, deque
+from collections import OrderedDict
 from copy import deepcopy
 from datetime import timedelta
-from typing import Any, Callable, Iterable, Mapping, MutableMapping, Optional
+from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional
 
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.declarative.incremental.global_substream_cursor import (
@@ -66,8 +66,8 @@ class ConcurrentPerPartitionCursor(Cursor):
     _GLOBAL_STATE_KEY = "state"
     _PERPARTITION_STATE_KEY = "states"
     _IS_PARTITION_DUPLICATION_LOGGED = False
-    _KEY = 0
-    _VALUE = 1
+    _PARENT_STATE = 0
+    _GENERATION_SEQUENCE = 1
 
     def __init__(
         self,
@@ -105,12 +105,12 @@ def __init__(
         self._parent_state: Optional[StreamState] = None
 
         # Tracks when the last slice for partition is emitted
-        self._finished_partitions: set[str] = set()
-        # Used to track the sequence numbers of open partitions
-        self._open_seqs: deque[int] = deque()
-        self._next_seq: int = 0
-        # Dictionary to map partition keys to their sequence numbers
-        self._seq_by_partition: dict[str, int] = {}
+        self._partitions_done_generating_stream_slices: set[str] = set()
+        # Used to track the index of partitions that are not closed yet
+        self._processing_partitions_indexes: List[int] = list()
+        self._generated_partitions_count: int = 0
+        # Dictionary to map partition keys to their index
+        self._partition_key_to_index: dict[str, int] = {}
 
         self._lock = threading.Lock()
         self._lookback_window: int = 0
@@ -167,7 +167,7 @@ def close_partition(self, partition: Partition) -> None:
                 self._cursor_per_partition[partition_key].close_partition(partition=partition)
                 cursor = self._cursor_per_partition[partition_key]
                 if (
-                    partition_key in self._finished_partitions
+                    partition_key in self._partitions_done_generating_stream_slices
                     and self._semaphore_per_partition[partition_key]._value == 0
                 ):
                     self._update_global_cursor(cursor.state[self.cursor_field.cursor_field_key])
@@ -188,7 +188,10 @@ def _check_and_update_parent_state(self) -> None:
             )
 
             # if any partition that started <= candidate_seq is still open, we must wait
-            if self._open_seqs and self._open_seqs[0] <= candidate_seq:
+            if (
+                self._processing_partitions_indexes
+                and self._processing_partitions_indexes[0] <= candidate_seq
+            ):
                 break
 
             # safe to pop
@@ -273,20 +276,21 @@ def _generate_slices_from_partition(
             if not self._IS_PARTITION_DUPLICATION_LOGGED:
                 logger.warning(f"Partition duplication detected for stream {self._stream_name}")
                 self._IS_PARTITION_DUPLICATION_LOGGED = True
+            return
         else:
             self._semaphore_per_partition[partition_key] = threading.Semaphore(0)
 
         with self._lock:
-            seq = self._next_seq
-            self._next_seq += 1
-            self._open_seqs.append(seq)
-            self._seq_by_partition[partition_key] = seq
+            seq = self._generated_partitions_count
+            self._generated_partitions_count += 1
+            self._processing_partitions_indexes.append(seq)
+            self._partition_key_to_index[partition_key] = seq
 
             if (
                 len(self._partition_parent_state_map) == 0
                 or self._partition_parent_state_map[
                     next(reversed(self._partition_parent_state_map))
-                ][0]
+                ][self._PARENT_STATE]
                 != parent_state
             ):
                 self._partition_parent_state_map[partition_key] = (deepcopy(parent_state), seq)
@@ -297,7 +301,7 @@ def _generate_slices_from_partition(
         ):
             self._semaphore_per_partition[partition_key].release()
             if is_last_slice:
-                self._finished_partitions.add(partition_key)
+                self._partitions_done_generating_stream_slices.add(partition_key)
             yield StreamSlice(
                 partition=partition, cursor_slice=cursor_slice, extra_fields=partition.extra_fields
             )
@@ -327,7 +331,7 @@ def _ensure_partition_limit(self) -> None:
             while len(self._cursor_per_partition) > self.DEFAULT_MAX_PARTITIONS_NUMBER - 1:
                 # Try removing finished partitions first
                 for partition_key in list(self._cursor_per_partition.keys()):
-                    if partition_key not in self._seq_by_partition:
+                    if partition_key not in self._partition_key_to_index:
                         oldest_partition = self._cursor_per_partition.pop(
                             partition_key
                         )  # Remove the oldest partition
@@ -466,16 +470,16 @@ def _cleanup_if_done(self, partition_key: str) -> None:
         cursor, semaphore, flag inside `_finished_partitions`
         """
         if not (
-            partition_key in self._finished_partitions
+            partition_key in self._partitions_done_generating_stream_slices
             and self._semaphore_per_partition[partition_key]._value == 0
         ):
             return
 
         self._semaphore_per_partition.pop(partition_key, None)
-        self._finished_partitions.discard(partition_key)
+        self._partitions_done_generating_stream_slices.discard(partition_key)
 
-        seq = self._seq_by_partition.pop(partition_key)
-        self._open_seqs.remove(seq)
+        seq = self._partition_key_to_index.pop(partition_key)
+        self._processing_partitions_indexes.remove(seq)
 
         logger.debug(f"Partition {partition_key} fully processed and cleaned up.")
 

unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py

@@ -3438,9 +3438,9 @@ def test_given_unfinished_first_parent_partition_no_parent_state_update():
     assert mock_cursor_2.stream_slices.call_count == 1  # Called once for each partition
 
     assert len(cursor._semaphore_per_partition) == 1
-    assert len(cursor._finished_partitions) == 1
-    assert len(cursor._open_seqs) == 1
-    assert len(cursor._seq_by_partition) == 1
+    assert len(cursor._partitions_done_generating_stream_slices) == 1
+    assert len(cursor._processing_partitions_indexes) == 1
+    assert len(cursor._partition_key_to_index) == 1
 
 
 def test_given_unfinished_last_parent_partition_with_partial_parent_state_update():
@@ -3526,9 +3526,9 @@ def test_given_unfinished_last_parent_partition_with_partial_parent_state_update
     assert mock_cursor_2.stream_slices.call_count == 1  # Called once for each partition
 
     assert len(cursor._semaphore_per_partition) == 1
-    assert len(cursor._finished_partitions) == 1
-    assert len(cursor._open_seqs) == 1
-    assert len(cursor._seq_by_partition) == 1
+    assert len(cursor._partitions_done_generating_stream_slices) == 1
+    assert len(cursor._processing_partitions_indexes) == 1
+    assert len(cursor._partition_key_to_index) == 1
 
 
 def test_given_all_partitions_finished_when_close_partition_then_final_state_emitted():
@@ -3606,9 +3606,9 @@ def test_given_all_partitions_finished_when_close_partition_then_final_state_emi
 
     # Checks that all internal variables are cleaned up
     assert len(cursor._semaphore_per_partition) == 0
-    assert len(cursor._finished_partitions) == 0
-    assert len(cursor._open_seqs) == 0
-    assert len(cursor._seq_by_partition) == 0
+    assert len(cursor._partitions_done_generating_stream_slices) == 0
+    assert len(cursor._processing_partitions_indexes) == 0
+    assert len(cursor._partition_key_to_index) == 0
 
 
 def test_given_partition_limit_exceeded_when_close_partition_then_switch_to_global_cursor():
@@ -3727,8 +3727,8 @@ def test_semaphore_cleanup():
     # Verify initial state
     assert len(cursor._semaphore_per_partition) == 2
     assert len(cursor._partition_parent_state_map) == 2
-    assert len(cursor._open_seqs) == 2
-    assert len(cursor._seq_by_partition) == 2
+    assert len(cursor._processing_partitions_indexes) == 2
+    assert len(cursor._partition_key_to_index) == 2
     assert cursor._partition_parent_state_map['{"id":"1"}'][0] == {"parent": {"state": "state1"}}
     assert cursor._partition_parent_state_map['{"id":"2"}'][0] == {"parent": {"state": "state2"}}
 
@@ -3737,10 +3737,10 @@ def test_semaphore_cleanup():
         cursor.close_partition(DeclarativePartition("test_stream", {}, MagicMock(), MagicMock(), s))
 
     # Check state after closing partitions
-    assert len(cursor._finished_partitions) == 0
+    assert len(cursor._partitions_done_generating_stream_slices) == 0
     assert len(cursor._semaphore_per_partition) == 0
-    assert len(cursor._open_seqs) == 0
-    assert len(cursor._seq_by_partition) == 0
+    assert len(cursor._processing_partitions_indexes) == 0
+    assert len(cursor._partition_key_to_index) == 0
     assert len(cursor._partition_parent_state_map) == 0  # All parent states should be popped
     assert cursor._parent_state == {"parent": {"state": "state2"}}  # Last parent state
 
@@ -3788,3 +3788,127 @@ def test_given_global_state_when_read_then_state_is_not_per_partition() -> None:
             "use_global_cursor": True,  # ensures that it is running the Concurrent CDK version as this is not populated in the declarative implementation
         },  # this state does have per partition which would be under `states`
     )
+
+
+def _make_inner_cursor(ts: str) -> MagicMock:
+    """Return an inner cursor that yields exactly one slice and has a proper state."""
+    inner = MagicMock()
+    inner.stream_slices.return_value = iter([{"dummy": "slice"}])
+    inner.state = {"updated_at": ts}
+    inner.close_partition.return_value = None
+    inner.observe.return_value = None
+    return inner
+
+
+def test_duplicate_partition_after_cleanup():
+    inner_cursors = [
+        _make_inner_cursor("2024-01-01T00:00:00Z"),  # for first "1"
+        _make_inner_cursor("2024-01-02T00:00:00Z"),  # for "2"
+        _make_inner_cursor("2024-01-03T00:00:00Z"),  # for second "1"
+    ]
+    cursor_factory_mock = MagicMock()
+    cursor_factory_mock.create.side_effect = inner_cursors
+
+    converter = CustomFormatConcurrentStreamStateConverter(
+        datetime_format="%Y-%m-%dT%H:%M:%SZ",
+        input_datetime_formats=["%Y-%m-%dT%H:%M:%SZ"],
+        is_sequential_state=True,
+        cursor_granularity=timedelta(0),
+    )
+
+    cursor = ConcurrentPerPartitionCursor(
+        cursor_factory=cursor_factory_mock,
+        partition_router=MagicMock(),
+        stream_name="dup_stream",
+        stream_namespace=None,
+        stream_state={},
+        message_repository=MagicMock(),
+        connector_state_manager=MagicMock(),
+        connector_state_converter=converter,
+        cursor_field=CursorField(cursor_field_key="updated_at"),
+    )
+
+    cursor.DEFAULT_MAX_PARTITIONS_NUMBER = 1
+
+    # ── Partition sequence: 1 → 2 → 1 ──────────────────────────────────
+    partitions = [
+        StreamSlice(partition={"id": "1"}, cursor_slice={}, extra_fields={}),
+        StreamSlice(partition={"id": "2"}, cursor_slice={}, extra_fields={}),
+        StreamSlice(partition={"id": "1"}, cursor_slice={}, extra_fields={}),
+    ]
+    pr = cursor._partition_router
+    pr.stream_slices.return_value = iter(partitions)
+    pr.get_stream_state.return_value = {}
+
+    # Iterate lazily so that the first "1" gets cleaned before
+    # the second "1" arrives.
+    slice_gen = cursor.stream_slices()
+
+    first_1 = next(slice_gen)
+    cursor.close_partition(
+        DeclarativePartition("dup_stream", {}, MagicMock(), MagicMock(), first_1)
+    )
+
+    two = next(slice_gen)
+    cursor.close_partition(DeclarativePartition("dup_stream", {}, MagicMock(), MagicMock(), two))
+
+    second_1 = next(slice_gen)
+    cursor.close_partition(
+        DeclarativePartition("dup_stream", {}, MagicMock(), MagicMock(), second_1)
+    )
+
+    assert cursor._IS_PARTITION_DUPLICATION_LOGGED is False  # No duplicate detected
+    assert len(cursor._semaphore_per_partition) == 0
+    assert len(cursor._processing_partitions_indexes) == 0
+    assert len(cursor._partition_key_to_index) == 0
+
+
+def test_duplicate_partition_while_processing():
+    inner_cursors = [
+        _make_inner_cursor("2024-01-01T00:00:00Z"),  # first “1”
+        _make_inner_cursor("2024-01-02T00:00:00Z"),  # “2”
+        _make_inner_cursor("2024-01-03T00:00:00Z"),  # for second "1"
+    ]
+
+    factory = MagicMock()
+    factory.create.side_effect = inner_cursors
+
+    cursor = ConcurrentPerPartitionCursor(
+        cursor_factory=factory,
+        partition_router=MagicMock(),
+        stream_name="dup_stream",
+        stream_namespace=None,
+        stream_state={},
+        message_repository=MagicMock(),
+        connector_state_manager=MagicMock(),
+        connector_state_converter=MagicMock(),
+        cursor_field=CursorField(cursor_field_key="updated_at"),
+    )
+
+    partitions = [
+        StreamSlice(partition={"id": "1"}, cursor_slice={}, extra_fields={}),
+        StreamSlice(partition={"id": "2"}, cursor_slice={}, extra_fields={}),
+        StreamSlice(partition={"id": "1"}, cursor_slice={}, extra_fields={}),
+    ]
+    pr = cursor._partition_router
+    pr.stream_slices.return_value = iter(partitions)
+    pr.get_stream_state.return_value = {}
+
+    generated = list(cursor.stream_slices())
+    # Only “1” and “2” emitted – duplicate “1” skipped
+    assert len(generated) == 2
+
+    # Close “2” first
+    cursor.close_partition(
+        DeclarativePartition("dup_stream", {}, MagicMock(), MagicMock(), generated[1])
+    )
+    # Now close the initial “1”
+    cursor.close_partition(
+        DeclarativePartition("dup_stream", {}, MagicMock(), MagicMock(), generated[0])
+    )
+
+    assert cursor._IS_PARTITION_DUPLICATION_LOGGED is True  # warning emitted
+    assert len(cursor._cursor_per_partition) == 2
+    assert len(cursor._semaphore_per_partition) == 0
+    assert len(cursor._processing_partitions_indexes) == 0
+    assert len(cursor._partition_key_to_index) == 0