airbyte-cdk 0.72.1__py3-none-any.whl → 6.13.1.dev4106__py3-none-any.whl

airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py

@@ -2,109 +2,167 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import json
-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.sources.declarative.incremental.cursor import Cursor
-from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer
-from airbyte_cdk.sources.declarative.types import Record, StreamSlice, StreamState
+from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
+from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
+from airbyte_cdk.sources.streams.checkpoint.per_partition_key_serializer import (
+    PerPartitionKeySerializer,
+)
+from airbyte_cdk.sources.types import Record, StreamSlice, StreamState
 
-
-class PerPartitionKeySerializer:
-    """
-    We are concerned of the performance of looping through the `states` list and evaluating equality on the partition. To reduce this
-    concern, we wanted to use dictionaries to map `partition -> cursor`. However, partitions are dict and dict can't be used as dict keys
-    since they are not hashable. By creating json string using the dict, we can have a use the dict as a key to the dict since strings are
-    hashable.
-    """
-
-    @staticmethod
-    def to_partition_key(to_serialize: Any) -> str:
-        # separators have changed in Python 3.4. To avoid being impacted by further change, we explicitly specify our own value
-        return json.dumps(to_serialize, indent=None, separators=(",", ":"), sort_keys=True)
-
-    @staticmethod
-    def to_partition(to_deserialize: Any) -> Mapping[str, Any]:
-        return json.loads(to_deserialize)  # type: ignore # The partition is known to be a dict, but the type hint is Any
+logger = logging.getLogger("airbyte")
 
 
 class CursorFactory:
-    def __init__(self, create_function: Callable[[], Cursor]):
+    def __init__(self, create_function: Callable[[], DeclarativeCursor]):
         self._create_function = create_function
 
-    def create(self) -> Cursor:
+    def create(self) -> DeclarativeCursor:
         return self._create_function()
 
 
-class PerPartitionCursor(Cursor):
+class PerPartitionCursor(DeclarativeCursor):
     """
-    Given a stream has many partitions, it is important to provide a state per partition.
-
-    Record | Stream Slice | Last Record | DatetimeCursorBased cursor
-    -- | -- | -- | --
-    1 | {"start_time": "2021-01-01","end_time": "2021-01-31","owner_resource": "1"''} | cursor_field: “2021-01-15” | 2021-01-15
-    2 | {"start_time": "2021-02-01","end_time": "2021-02-28","owner_resource": "1"''} | cursor_field: “2021-02-15” | 2021-02-15
-    3 | {"start_time": "2021-01-01","end_time": "2021-01-31","owner_resource": "2"''} | cursor_field: “2021-01-03” | 2021-01-03
-    4 | {"start_time": "2021-02-01","end_time": "2021-02-28","owner_resource": "2"''} | cursor_field: “2021-02-14” | 2021-02-14
-
-    Given the following errors, this can lead to some loss or duplication of records:
-    When | Problem | Affected Record
-    -- | -- | --
-    Between record #1 and #2 | Loss | #3
-    Between record #2 and #3 | Loss | #3, #4
-    Between record #3 and #4 | Duplication | #1, #2
-
-    Therefore, we need to manage state per partition.
+    Manages state per partition when a stream has many partitions, to prevent data loss or duplication.
+
+    **Partition Limitation and Limit Reached Logic**
+
+    - **DEFAULT_MAX_PARTITIONS_NUMBER**: The maximum number of partitions to keep in memory (default is 10,000).
+    - **_cursor_per_partition**: An ordered dictionary that stores cursors for each partition.
+    - **_over_limit**: A counter that increments each time an oldest partition is removed when the limit is exceeded.
+
+    The class ensures that the number of partitions tracked does not exceed the `DEFAULT_MAX_PARTITIONS_NUMBER` to prevent excessive memory usage.
+
+    - When the number of partitions exceeds the limit, the oldest partitions are removed from `_cursor_per_partition`, and `_over_limit` is incremented accordingly.
+    - The `limit_reached` method returns `True` when `_over_limit` exceeds `DEFAULT_MAX_PARTITIONS_NUMBER`, indicating that the global cursor should be used instead of per-partition cursors.
+
+    This approach avoids unnecessary switching to a global cursor due to temporary spikes in partition counts, ensuring that switching is only done when a sustained high number of partitions is observed.
     """
 
+    DEFAULT_MAX_PARTITIONS_NUMBER = 10000
     _NO_STATE: Mapping[str, Any] = {}
     _NO_CURSOR_STATE: Mapping[str, Any] = {}
     _KEY = 0
     _VALUE = 1
+    _state_to_migrate_from: Mapping[str, Any] = {}
 
-    def __init__(self, cursor_factory: CursorFactory, partition_router: StreamSlicer):
+    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, Cursor] = {}
+        # 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._over_limit = 0
         self._partition_serializer = PerPartitionKeySerializer()
 
     def stream_slices(self) -> Iterable[StreamSlice]:
         slices = self._partition_router.stream_slices()
         for partition in slices:
-            cursor = self._cursor_per_partition.get(self._to_partition_key(partition.partition))
-            if not cursor:
-                cursor = self._create_cursor(self._NO_CURSOR_STATE)
-                self._cursor_per_partition[self._to_partition_key(partition.partition)] = cursor
+            yield from self.generate_slices_from_partition(partition)
+
+    def generate_slices_from_partition(self, partition: StreamSlice) -> Iterable[StreamSlice]:
+        # 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:
+            partition_state = (
+                self._state_to_migrate_from
+                if self._state_to_migrate_from
+                else self._NO_CURSOR_STATE
+            )
+            cursor = self._create_cursor(partition_state)
+            self._cursor_per_partition[self._to_partition_key(partition.partition)] = cursor
+
+        for cursor_slice in cursor.stream_slices():
+            yield StreamSlice(
+                partition=partition, cursor_slice=cursor_slice, extra_fields=partition.extra_fields
+            )
+
+    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:
+            self._over_limit += 1
+            oldest_partition = self._cursor_per_partition.popitem(last=False)[
+                0
+            ]  # Remove the oldest partition
+            logger.warning(
+                f"The maximum number of partitions has been reached. Dropping the oldest partition: {oldest_partition}. Over limit: {self._over_limit}."
+            )
 
-            for cursor_slice in cursor.stream_slices():
-                yield StreamSlice(partition=partition, cursor_slice=cursor_slice)
+    def limit_reached(self) -> bool:
+        return self._over_limit > self.DEFAULT_MAX_PARTITIONS_NUMBER
 
     def set_initial_state(self, stream_state: StreamState) -> None:
+        """
+        Set the initial state for the cursors.
+
+        This method initializes the state for each partition cursor using the provided stream state.
+        If a partition state is provided in the stream state, it will update the corresponding partition cursor with this state.
+
+        Additionally, it sets the parent state for partition routers that are based on parent streams. If a partition router
+        does not have parent streams, this step will be skipped due to the default PartitionRouter implementation.
+
+        Args:
+            stream_state (StreamState): The state of the streams to be set. The format of the stream state should be:
+                {
+                    "states": [
+                        {
+                            "partition": {
+                                "partition_key": "value"
+                            },
+                            "cursor": {
+                                "last_updated": "2023-05-27T00:00:00Z"
+                            }
+                        }
+                    ],
+                    "parent_state": {
+                        "parent_stream_name": {
+                            "last_updated": "2023-05-27T00:00:00Z"
+                        }
+                    }
+                }
+        """
         if not stream_state:
             return
 
-        for state in stream_state["states"]:
-            self._cursor_per_partition[self._to_partition_key(state["partition"])] = self._create_cursor(state["cursor"])
+        if "states" not in stream_state:
+            # We assume that `stream_state` is in a global format that can be applied to all partitions.
+            # Example: {"global_state_format_key": "global_state_format_value"}
+            self._state_to_migrate_from = stream_state
+
+        else:
+            for state in stream_state["states"]:
+                self._cursor_per_partition[self._to_partition_key(state["partition"])] = (
+                    self._create_cursor(state["cursor"])
+                )
+
+            # set default state for missing partitions if it is per partition with fallback to global
+            if "state" in stream_state:
+                self._state_to_migrate_from = stream_state["state"]
+
+        # Set parent state for partition routers based on parent streams
+        self._partition_router.set_initial_state(stream_state)
 
     def observe(self, stream_slice: StreamSlice, record: Record) -> None:
         self._cursor_per_partition[self._to_partition_key(stream_slice.partition)].observe(
             StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice), record
         )
 
-    def close_slice(self, stream_slice: StreamSlice, most_recent_record: Optional[Record]) -> None:
+    def close_slice(self, stream_slice: StreamSlice, *args: Any) -> None:
         try:
-            cursor_most_recent_record = (
-                Record(most_recent_record.data, StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice))
-                if most_recent_record
-                else most_recent_record
-            )
             self._cursor_per_partition[self._to_partition_key(stream_slice.partition)].close_slice(
-                StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice), cursor_most_recent_record
+                StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice), *args
             )
         except KeyError as exception:
             raise ValueError(
                 f"Partition {str(exception)} could not be found in current state based on the record. This is unexpected because "
-                f"we should only update state for partition that where emitted during `stream_slices`"
+                f"we should only update state for partitions that were emitted during `stream_slices`"
             )
 
     def get_stream_state(self) -> StreamState:
@@ -118,7 +176,12 @@ class PerPartitionCursor(Cursor):
                         "cursor": cursor_state,
                     }
                 )
-        return {"states": states}
+        state: dict[str, Any] = {"states": states}
+
+        parent_state = self._partition_router.get_stream_state()
+        if parent_state:
+            state["parent_state"] = parent_state
+        return state
 
     def _get_state_for_partition(self, partition: Mapping[str, Any]) -> Optional[StreamState]:
         cursor = self._cursor_per_partition.get(self._to_partition_key(partition))
@@ -146,7 +209,7 @@ class PerPartitionCursor(Cursor):
 
         return self._get_state_for_partition(stream_slice.partition)
 
-    def _create_cursor(self, cursor_state: Any) -> Cursor:
+    def _create_cursor(self, cursor_state: Any) -> DeclarativeCursor:
         cursor = self._cursor_factory.create()
         cursor.set_initial_state(cursor_state)
         return cursor
@@ -163,7 +226,9 @@ class PerPartitionCursor(Cursor):
                 stream_state=stream_state,
                 stream_slice=StreamSlice(partition=stream_slice.partition, cursor_slice={}),
                 next_page_token=next_page_token,
-            ) | self._cursor_per_partition[self._to_partition_key(stream_slice.partition)].get_request_params(
+            ) | self._cursor_per_partition[
+                self._to_partition_key(stream_slice.partition)
+            ].get_request_params(
                 stream_state=stream_state,
                 stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice),
                 next_page_token=next_page_token,
@@ -183,7 +248,9 @@ class PerPartitionCursor(Cursor):
                 stream_state=stream_state,
                 stream_slice=StreamSlice(partition=stream_slice.partition, cursor_slice={}),
                 next_page_token=next_page_token,
-            ) | self._cursor_per_partition[self._to_partition_key(stream_slice.partition)].get_request_headers(
+            ) | self._cursor_per_partition[
+                self._to_partition_key(stream_slice.partition)
+            ].get_request_headers(
                 stream_state=stream_state,
                 stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice),
                 next_page_token=next_page_token,
@@ -203,7 +270,9 @@ class PerPartitionCursor(Cursor):
                 stream_state=stream_state,
                 stream_slice=StreamSlice(partition=stream_slice.partition, cursor_slice={}),
                 next_page_token=next_page_token,
-            ) | self._cursor_per_partition[self._to_partition_key(stream_slice.partition)].get_request_body_data(
+            ) | self._cursor_per_partition[
+                self._to_partition_key(stream_slice.partition)
+            ].get_request_body_data(
                 stream_state=stream_state,
                 stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice),
                 next_page_token=next_page_token,
@@ -223,7 +292,9 @@ class PerPartitionCursor(Cursor):
                 stream_state=stream_state,
                 stream_slice=StreamSlice(partition=stream_slice.partition, cursor_slice={}),
                 next_page_token=next_page_token,
-            ) | self._cursor_per_partition[self._to_partition_key(stream_slice.partition)].get_request_body_json(
+            ) | self._cursor_per_partition[
+                self._to_partition_key(stream_slice.partition)
+            ].get_request_body_json(
                 stream_state=stream_state,
                 stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice),
                 next_page_token=next_page_token,
@@ -232,32 +303,46 @@ class PerPartitionCursor(Cursor):
             raise ValueError("A partition needs to be provided in order to get request body json")
 
     def should_be_synced(self, record: Record) -> bool:
-        return self._get_cursor(record).should_be_synced(self._convert_record_to_cursor_record(record))
+        return self._get_cursor(record).should_be_synced(
+            self._convert_record_to_cursor_record(record)
+        )
 
     def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
         if not first.associated_slice or not second.associated_slice:
-            raise ValueError(f"Both records should have an associated slice but got {first.associated_slice} and {second.associated_slice}")
+            raise ValueError(
+                f"Both records should have an associated slice but got {first.associated_slice} and {second.associated_slice}"
+            )
         if first.associated_slice.partition != second.associated_slice.partition:
             raise ValueError(
                 f"To compare records, partition should be the same but got {first.associated_slice.partition} and {second.associated_slice.partition}"
             )
 
         return self._get_cursor(first).is_greater_than_or_equal(
-            self._convert_record_to_cursor_record(first), self._convert_record_to_cursor_record(second)
+            self._convert_record_to_cursor_record(first),
+            self._convert_record_to_cursor_record(second),
         )
 
     @staticmethod
     def _convert_record_to_cursor_record(record: Record) -> Record:
         return Record(
-            record.data,
-            StreamSlice(partition={}, cursor_slice=record.associated_slice.cursor_slice) if record.associated_slice else None,
+            data=record.data,
+            stream_name=record.stream_name,
+            associated_slice=StreamSlice(
+                partition={}, cursor_slice=record.associated_slice.cursor_slice
+            )
+            if record.associated_slice
+            else None,
         )
 
-    def _get_cursor(self, record: Record) -> Cursor:
+    def _get_cursor(self, record: Record) -> DeclarativeCursor:
         if not record.associated_slice:
-            raise ValueError("Invalid state as stream slices that are emitted should refer to an existing cursor")
+            raise ValueError(
+                "Invalid state as stream slices that are emitted should refer to an existing cursor"
+            )
         partition_key = self._to_partition_key(record.associated_slice.partition)
         if partition_key not in self._cursor_per_partition:
-            raise ValueError("Invalid state as stream slices that are emitted should refer to an existing cursor")
+            raise ValueError(
+                "Invalid state as stream slices that are emitted should refer to an existing cursor"
+            )
         cursor = self._cursor_per_partition[partition_key]
         return cursor

airbyte_cdk/sources/declarative/incremental/per_partition_with_global.py

@@ -0,0 +1,200 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+from typing import Any, Iterable, Mapping, MutableMapping, Optional, Union
+
+from airbyte_cdk.sources.declarative.incremental.datetime_based_cursor import DatetimeBasedCursor
+from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
+from airbyte_cdk.sources.declarative.incremental.global_substream_cursor import (
+    GlobalSubstreamCursor,
+    iterate_with_last_flag_and_state,
+)
+from airbyte_cdk.sources.declarative.incremental.per_partition_cursor import (
+    CursorFactory,
+    PerPartitionCursor,
+)
+from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter
+from airbyte_cdk.sources.types import Record, StreamSlice, StreamState
+
+
+class PerPartitionWithGlobalCursor(DeclarativeCursor):
+    """
+    Manages state for streams with multiple partitions, with an optional fallback to a global cursor when specific conditions are met.
+
+    This cursor handles partitioned streams by maintaining individual state per partition using `PerPartitionCursor`. If the number of partitions exceeds a defined limit, it switches to a global cursor (`GlobalSubstreamCursor`) to manage state more efficiently.
+
+    **Overview**
+
+    - **Partition-Based State**: Initially manages state per partition to ensure accurate processing of each partition's data.
+    - **Global Fallback**: Switches to a global cursor when the partition limit is exceeded to handle state management more effectively.
+
+    **Switching Logic**
+
+    - Monitors the number of partitions.
+    - If `PerPartitionCursor.limit_reached()` returns `True`, sets `_use_global_cursor` to `True`, activating the global cursor.
+
+    **Active Cursor Selection**
+
+    - Uses the `_get_active_cursor()` helper method to select the active cursor based on the `_use_global_cursor` flag.
+    - This simplifies the logic and ensures consistent cursor usage across methods.
+
+    **State Structure Example**
+
+    ```json
+    {
+        "states": [
+            {
+                "partition": {"partition_key": "partition_1"},
+                "cursor": {"cursor_field": "2021-01-15"}
+            },
+            {
+                "partition": {"partition_key": "partition_2"},
+                "cursor": {"cursor_field": "2021-02-14"}
+            }
+        ],
+        "state": {
+            "cursor_field": "2021-02-15"
+        },
+        "use_global_cursor": false
+    }
+    ```
+
+    In this example, the cursor is using partition-based state management (`"use_global_cursor": false`), maintaining separate cursor states for each partition.
+
+    **Usage Scenario**
+
+    Suitable for streams where the number of partitions may vary significantly, requiring dynamic switching between per-partition and global state management to ensure data consistency and efficient synchronization.
+    """
+
+    def __init__(
+        self,
+        cursor_factory: CursorFactory,
+        partition_router: PartitionRouter,
+        stream_cursor: DatetimeBasedCursor,
+    ):
+        self._partition_router = partition_router
+        self._per_partition_cursor = PerPartitionCursor(cursor_factory, partition_router)
+        self._global_cursor = GlobalSubstreamCursor(stream_cursor, partition_router)
+        self._use_global_cursor = False
+        self._current_partition: Optional[Mapping[str, Any]] = None
+        self._last_slice: bool = False
+        self._parent_state: Optional[Mapping[str, Any]] = None
+
+    def _get_active_cursor(self) -> Union[PerPartitionCursor, GlobalSubstreamCursor]:
+        return self._global_cursor if self._use_global_cursor else self._per_partition_cursor
+
+    def stream_slices(self) -> Iterable[StreamSlice]:
+        self._global_cursor.start_slices_generation()
+
+        # Iterate through partitions and process slices
+        for partition, is_last_partition, parent_state in iterate_with_last_flag_and_state(
+            self._partition_router.stream_slices(), self._partition_router.get_stream_state
+        ):
+            # Generate slices for the current cursor and handle the last slice using the flag
+            self._parent_state = parent_state
+            for slice, is_last_slice, _ in iterate_with_last_flag_and_state(
+                self._get_active_cursor().generate_slices_from_partition(partition=partition),
+                lambda: None,
+            ):
+                self._global_cursor.register_slice(is_last_slice and is_last_partition)
+                yield slice
+        self._parent_state = self._partition_router.get_stream_state()
+
+    def set_initial_state(self, stream_state: StreamState) -> None:
+        """
+        Set the initial state for the cursors.
+        """
+        self._use_global_cursor = stream_state.get("use_global_cursor", False)
+
+        self._parent_state = stream_state.get("parent_state", {})
+
+        self._global_cursor.set_initial_state(stream_state)
+        if not self._use_global_cursor:
+            self._per_partition_cursor.set_initial_state(stream_state)
+
+    def observe(self, stream_slice: StreamSlice, record: Record) -> None:
+        if not self._use_global_cursor and self._per_partition_cursor.limit_reached():
+            self._use_global_cursor = True
+
+        if not self._use_global_cursor:
+            self._per_partition_cursor.observe(stream_slice, record)
+        self._global_cursor.observe(stream_slice, record)
+
+    def close_slice(self, stream_slice: StreamSlice, *args: Any) -> None:
+        if not self._use_global_cursor:
+            self._per_partition_cursor.close_slice(stream_slice, *args)
+        self._global_cursor.close_slice(stream_slice, *args)
+
+    def get_stream_state(self) -> StreamState:
+        final_state: MutableMapping[str, Any] = {"use_global_cursor": self._use_global_cursor}
+
+        final_state.update(self._global_cursor.get_stream_state())
+        if not self._use_global_cursor:
+            final_state.update(self._per_partition_cursor.get_stream_state())
+
+        final_state["parent_state"] = self._parent_state
+        if not final_state.get("parent_state"):
+            del final_state["parent_state"]
+
+        return final_state
+
+    def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
+        return self._get_active_cursor().select_state(stream_slice)
+
+    def get_request_params(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return self._get_active_cursor().get_request_params(
+            stream_state=stream_state,
+            stream_slice=stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def get_request_headers(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return self._get_active_cursor().get_request_headers(
+            stream_state=stream_state,
+            stream_slice=stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def get_request_body_data(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Union[Mapping[str, Any], str]:
+        return self._get_active_cursor().get_request_body_data(
+            stream_state=stream_state,
+            stream_slice=stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def get_request_body_json(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return self._get_active_cursor().get_request_body_json(
+            stream_state=stream_state,
+            stream_slice=stream_slice,
+            next_page_token=next_page_token,
+        )
+
+    def should_be_synced(self, record: Record) -> bool:
+        return self._get_active_cursor().should_be_synced(record)
+
+    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
+        return self._global_cursor.is_greater_than_or_equal(first, second)

airbyte_cdk/sources/declarative/incremental/resumable_full_refresh_cursor.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from dataclasses import InitVar, dataclass
+from typing import Any, Iterable, Mapping, Optional
+
+from airbyte_cdk.sources.declarative.incremental import DeclarativeCursor
+from airbyte_cdk.sources.declarative.types import Record, StreamSlice, StreamState
+from airbyte_cdk.sources.streams.checkpoint.checkpoint_reader import FULL_REFRESH_COMPLETE_STATE
+
+
+@dataclass
+class ResumableFullRefreshCursor(DeclarativeCursor):
+    parameters: InitVar[Mapping[str, Any]]
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        self._cursor: StreamState = {}
+
+    def get_stream_state(self) -> StreamState:
+        return self._cursor
+
+    def set_initial_state(self, stream_state: StreamState) -> None:
+        self._cursor = stream_state
+
+    def observe(self, stream_slice: StreamSlice, record: Record) -> None:
+        """
+        Resumable full refresh manages state using a page number so it does not need to update state by observing incoming records.
+        """
+        pass
+
+    def close_slice(self, stream_slice: StreamSlice, *args: Any) -> None:
+        # The ResumableFullRefreshCursor doesn't support nested streams yet so receiving a partition is unexpected
+        if stream_slice.partition:
+            raise ValueError(
+                f"Stream slice {stream_slice} should not have a partition. Got {stream_slice.partition}."
+            )
+        self._cursor = stream_slice.cursor_slice
+
+    def should_be_synced(self, record: Record) -> bool:
+        """
+        Unlike date-based cursors which filter out records outside slice boundaries, resumable full refresh records exist within pages
+        that don't have filterable bounds. We should always return them.
+        """
+        return True
+
+    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
+        """
+        RFR record don't have ordering to be compared between one another.
+        """
+        return False
+
+    def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
+        # A top-level RFR cursor only manages the state of a single partition
+        return self._cursor
+
+    def stream_slices(self) -> Iterable[StreamSlice]:
+        """
+        Resumable full refresh cursors only return a single slice and can't perform partitioning because iteration is done per-page
+        along an unbounded set.
+        """
+        yield from [StreamSlice(cursor_slice=self._cursor, partition={})]
+
+    # This is an interesting pattern that might not seem obvious at first glance. This cursor itself has no functional need to
+    # inject any request values into the outbound response because the up-to-date pagination state is already loaded and
+    # maintained by the paginator component
+    def get_request_params(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return {}
+
+    def get_request_headers(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return {}
+
+    def get_request_body_data(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return {}
+
+    def get_request_body_json(
+        self,
+        *,
+        stream_state: Optional[StreamState] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        next_page_token: Optional[Mapping[str, Any]] = None,
+    ) -> Mapping[str, Any]:
+        return {}
+
+
+@dataclass
+class ChildPartitionResumableFullRefreshCursor(ResumableFullRefreshCursor):
+    """
+    The Sub-stream Resumable Cursor for Full-Refresh substreams.
+    Follows the parent type `ResumableFullRefreshCursor` with a small override,
+    to provide the ability to close the substream's slice once it has finished processing.
+
+    Check the `close_slice` method overide for more info about the actual behaviour of this cursor.
+    """
+
+    def close_slice(self, stream_slice: StreamSlice, *args: Any) -> None:
+        """
+        Once the current slice has finished syncing:
+         - paginator returns None
+         - no more slices to process
+
+        we assume that the records are processed and emitted already,
+        thus we have to set the cursor to ` __ab_full_refresh_sync_complete: true `,
+        otherwise there is a risk of Inf. Loop processing the same slice.
+        """
+        self._cursor = FULL_REFRESH_COMPLETE_STATE