airbyte-cdk 0.0.0.dev0__py3-none-any.whl

airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py

@@ -0,0 +1,380 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+from collections import OrderedDict
+from typing import Any, Callable, Iterable, Mapping, Optional, Union
+
+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
+
+logger = logging.getLogger("airbyte")
+
+
+class CursorFactory:
+    def __init__(self, create_function: Callable[[], DeclarativeCursor]):
+        self._create_function = create_function
+
+    def create(self) -> DeclarativeCursor:
+        return self._create_function()
+
+
+class PerPartitionCursor(DeclarativeCursor):
+    """
+    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: PartitionRouter):
+        self._cursor_factory = cursor_factory
+        self._partition_router = partition_router
+        # 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:
+            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}."
+            )
+
+    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
+
+        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, *args: Any) -> None:
+        try:
+            self._cursor_per_partition[self._to_partition_key(stream_slice.partition)].close_slice(
+                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 partitions that were emitted during `stream_slices`"
+            )
+
+    def get_stream_state(self) -> StreamState:
+        states = []
+        for partition_tuple, cursor in self._cursor_per_partition.items():
+            cursor_state = cursor.get_stream_state()
+            if cursor_state:
+                states.append(
+                    {
+                        "partition": self._to_dict(partition_tuple),
+                        "cursor": cursor_state,
+                    }
+                )
+        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))
+        if cursor:
+            return cursor.get_stream_state()
+
+        return None
+
+    @staticmethod
+    def _is_new_state(stream_state: Mapping[str, Any]) -> bool:
+        return not bool(stream_state)
+
+    def _to_partition_key(self, partition: Mapping[str, Any]) -> str:
+        return self._partition_serializer.to_partition_key(partition)
+
+    def _to_dict(self, partition_key: str) -> Mapping[str, Any]:
+        return self._partition_serializer.to_partition(partition_key)
+
+    def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
+        if not stream_slice:
+            raise ValueError("A partition needs to be provided in order to extract a state")
+
+        if not stream_slice:
+            return None
+
+        return self._get_state_for_partition(stream_slice.partition)
+
+    def _create_cursor(self, cursor_state: Any) -> DeclarativeCursor:
+        cursor = self._cursor_factory.create()
+        cursor.set_initial_state(cursor_state)
+        return cursor
+
+    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]:
+        if stream_slice:
+            if self._to_partition_key(stream_slice.partition) not in self._cursor_per_partition:
+                self._create_cursor_for_partition(self._to_partition_key(stream_slice.partition))
+            return self._partition_router.get_request_params(  # type: ignore # this always returns a mapping
+                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(
+                stream_state=stream_state,
+                stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice),
+                next_page_token=next_page_token,
+            )
+        else:
+            raise ValueError("A partition needs to be provided in order to get request params")
+
+    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]:
+        if stream_slice:
+            if self._to_partition_key(stream_slice.partition) not in self._cursor_per_partition:
+                self._create_cursor_for_partition(self._to_partition_key(stream_slice.partition))
+            return self._partition_router.get_request_headers(  # type: ignore # this always returns a mapping
+                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(
+                stream_state=stream_state,
+                stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice),
+                next_page_token=next_page_token,
+            )
+        else:
+            raise ValueError("A partition needs to be provided in order to get request headers")
+
+    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]:
+        if stream_slice:
+            if self._to_partition_key(stream_slice.partition) not in self._cursor_per_partition:
+                self._create_cursor_for_partition(self._to_partition_key(stream_slice.partition))
+            return self._partition_router.get_request_body_data(  # type: ignore # this always returns a mapping
+                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(
+                stream_state=stream_state,
+                stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice),
+                next_page_token=next_page_token,
+            )
+        else:
+            raise ValueError("A partition needs to be provided in order to get request body data")
+
+    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]:
+        if stream_slice:
+            if self._to_partition_key(stream_slice.partition) not in self._cursor_per_partition:
+                self._create_cursor_for_partition(self._to_partition_key(stream_slice.partition))
+            return self._partition_router.get_request_body_json(  # type: ignore # this always returns a mapping
+                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(
+                stream_state=stream_state,
+                stream_slice=StreamSlice(partition={}, cursor_slice=stream_slice.cursor_slice),
+                next_page_token=next_page_token,
+            )
+        else:
+            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)
+        )
+
+    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}"
+            )
+        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),
+        )
+
+    @staticmethod
+    def _convert_record_to_cursor_record(record: Record) -> Record:
+        return Record(
+            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) -> DeclarativeCursor:
+        if not record.associated_slice:
+            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:
+            self._create_cursor_for_partition(partition_key)
+        cursor = self._cursor_per_partition[partition_key]
+        return cursor
+
+    def _create_cursor_for_partition(self, partition_key: str) -> None:
+        """
+        Dynamically creates and initializes a cursor for the specified partition.
+
+        This method is required for `ConcurrentPerPartitionCursor`. For concurrent cursors,
+        stream_slices is executed only for the concurrent cursor, so cursors per partition
+        are not created for the declarative cursor. This method ensures that a cursor is available
+        to create requests for the specified partition. The cursor is initialized
+        with the per-partition state if present in the initial state, or with the global state
+        adjusted by the lookback window, or with the state to migrate from.
+
+        Note:
+            This is a temporary workaround and should be removed once the declarative cursor
+            is decoupled from the concurrent cursor implementation.
+
+        Args:
+            partition_key (str): The unique identifier for the partition for which the cursor
+            needs to be created.
+        """
+        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[partition_key] = 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)