feat(airbyte-cdk): Add limitation for number of partitions to PerPartitionCursor (#42406)

Author: tolik0

airbyte-cdk/python/airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py

@@ -2,7 +2,9 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from typing import Any, Callable, Iterable, Mapping, MutableMapping, Optional, Union
+import logging
+from collections import OrderedDict
+from typing import Any, Callable, Iterable, Mapping, Optional, Union
 
 from airbyte_cdk.models import FailureType
 from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
@@ -41,6 +43,7 @@ class PerPartitionCursor(DeclarativeCursor):
     Therefore, we need to manage state per partition.
     """
 
+    DEFAULT_MAX_PARTITIONS_NUMBER = 10000
     _NO_STATE: Mapping[str, Any] = {}
     _NO_CURSOR_STATE: Mapping[str, Any] = {}
     _KEY = 0
@@ -49,12 +52,17 @@ class PerPartitionCursor(DeclarativeCursor):
     def __init__(self, cursor_factory: CursorFactory, partition_router: PartitionRouter):
         self._cursor_factory = cursor_factory
         self._partition_router = partition_router
-        self._cursor_per_partition: MutableMapping[str, DeclarativeCursor] = {}
+        # The dict is ordered to ensure that once the maximum number of partitions is reached,
+        # the oldest partitions can be efficiently removed, maintaining the most recent partitions.
+        self._cursor_per_partition: OrderedDict[str, DeclarativeCursor] = OrderedDict()
         self._partition_serializer = PerPartitionKeySerializer()
 
     def stream_slices(self) -> Iterable[StreamSlice]:
         slices = self._partition_router.stream_slices()
         for partition in slices:
+            # Ensure the maximum number of partitions is not exceeded
+            self._ensure_partition_limit()
+
             cursor = self._cursor_per_partition.get(self._to_partition_key(partition.partition))
             if not cursor:
                 cursor = self._create_cursor(self._NO_CURSOR_STATE)
@@ -63,6 +71,14 @@ def stream_slices(self) -> Iterable[StreamSlice]:
             for cursor_slice in cursor.stream_slices():
                 yield StreamSlice(partition=partition, cursor_slice=cursor_slice)
 
+    def _ensure_partition_limit(self) -> None:
+        """
+        Ensure the maximum number of partitions is not exceeded. If so, the oldest added partition will be dropped.
+        """
+        while len(self._cursor_per_partition) > self.DEFAULT_MAX_PARTITIONS_NUMBER - 1:
+            oldest_partition = self._cursor_per_partition.popitem(last=False)[0]  # Remove the oldest partition
+            logging.warning(f"The maximum number of partitions has been reached. Dropping the oldest partition: {oldest_partition}.")
+
     def set_initial_state(self, stream_state: StreamState) -> None:
         """
         Set the initial state for the cursors.

airbyte-cdk/python/unit_tests/sources/declarative/incremental/test_per_partition_cursor_integration.py

@@ -2,10 +2,21 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
-from airbyte_cdk.models import SyncMode
-from airbyte_cdk.sources.declarative.incremental.per_partition_cursor import StreamSlice
+from airbyte_cdk.models import (
+    AirbyteStateBlob,
+    AirbyteStateMessage,
+    AirbyteStateType,
+    AirbyteStream,
+    AirbyteStreamState,
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteStream,
+    DestinationSyncMode,
+    StreamDescriptor,
+    SyncMode,
+)
+from airbyte_cdk.sources.declarative.incremental.per_partition_cursor import PerPartitionCursor, StreamSlice
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever
 from airbyte_cdk.sources.types import Record
@@ -268,3 +279,93 @@ def test_substream_without_input_state():
                 cursor_slice={"start_time": "2022-02-01", "end_time": "2022-02-28"},
             ),
         ]
+
+
+def test_partition_limitation():
+    source = ManifestDeclarativeSource(
+        source_config=ManifestBuilder()
+        .with_list_partition_router("Rates", "partition_field", ["1", "2", "3"])
+        .with_incremental_sync(
+            "Rates",
+            start_datetime="2022-01-01",
+            end_datetime="2022-02-28",
+            datetime_format="%Y-%m-%d",
+            cursor_field=CURSOR_FIELD,
+            step="P1M",
+            cursor_granularity="P1D",
+        )
+        .build()
+    )
+
+    partition_slices = [
+        StreamSlice(partition={"partition_field": "1"}, cursor_slice={}),
+        StreamSlice(partition={"partition_field": "2"}, cursor_slice={}),
+        StreamSlice(partition={"partition_field": "3"}, cursor_slice={}),
+    ]
+
+    records_list = [
+        [
+            Record({"a record key": "a record value", CURSOR_FIELD: "2022-01-15"}, partition_slices[0]),
+            Record({"a record key": "a record value", CURSOR_FIELD: "2022-01-16"}, partition_slices[0]),
+        ],
+        [Record({"a record key": "a record value", CURSOR_FIELD: "2022-02-15"}, partition_slices[0])],
+        [Record({"a record key": "a record value", CURSOR_FIELD: "2022-01-16"}, partition_slices[1])],
+        [],
+        [],
+        [Record({"a record key": "a record value", CURSOR_FIELD: "2022-02-17"}, partition_slices[2])],
+    ]
+
+    configured_stream = ConfiguredAirbyteStream(
+        stream=AirbyteStream(name="Rates", json_schema={}, supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental]),
+        sync_mode=SyncMode.incremental,
+        destination_sync_mode=DestinationSyncMode.append,
+    )
+    catalog = ConfiguredAirbyteCatalog(streams=[configured_stream])
+
+    initial_state = [
+        AirbyteStateMessage(
+            type=AirbyteStateType.STREAM,
+            stream=AirbyteStreamState(
+                stream_descriptor=StreamDescriptor(name="post_comment_votes", namespace=None),
+                stream_state=AirbyteStateBlob.parse_obj(
+                    {
+                        "states": [
+                            {
+                                "partition": {"partition_field": "1"},
+                                "cursor": {CURSOR_FIELD: "2022-01-01"},
+                            },
+                            {
+                                "partition": {"partition_field": "2"},
+                                "cursor": {CURSOR_FIELD: "2022-01-02"},
+                            },
+                            {
+                                "partition": {"partition_field": "3"},
+                                "cursor": {CURSOR_FIELD: "2022-01-03"},
+                            },
+                        ]
+                    }
+                ),
+            ),
+        )
+    ]
+    logger = MagicMock()
+
+    # with patch.object(PerPartitionCursor, "stream_slices", return_value=partition_slices):
+    with patch.object(SimpleRetriever, "_read_pages", side_effect=records_list):
+        with patch.object(PerPartitionCursor, "DEFAULT_MAX_PARTITIONS_NUMBER", 2):
+            output = list(source.read(logger, {}, catalog, initial_state))
+
+    # assert output_data == expected_records
+    final_state = [message.state.stream.stream_state.dict() for message in output if message.state]
+    assert final_state[-1] == {
+        "states": [
+            {
+                "partition": {"partition_field": "2"},
+                "cursor": {CURSOR_FIELD: "2022-01-16"},
+            },
+            {
+                "partition": {"partition_field": "3"},
+                "cursor": {CURSOR_FIELD: "2022-02-17"},
+            },
+        ]
+    }