Use request options provider for ConcurrentPerPartitionCursor

Author: tolik0

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -3,8 +3,6 @@
 #
 
 import logging
-from airbyte_cdk.sources.declarative.incremental.per_partition_cursor import PerPartitionCursor
-from airbyte_cdk.sources.declarative.incremental.per_partition_with_global import PerPartitionWithGlobalCursor
 from typing import Any, Callable, Generic, Iterator, List, Mapping, Optional, Tuple, Union
 
 from airbyte_cdk.models import (
@@ -22,6 +20,9 @@
     ClientSideIncrementalRecordFilterDecorator,
 )
 from airbyte_cdk.sources.declarative.incremental.datetime_based_cursor import DatetimeBasedCursor
+from airbyte_cdk.sources.declarative.incremental.per_partition_with_global import (
+    PerPartitionWithGlobalCursor,
+)
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
@@ -306,8 +307,8 @@ def _group_streams(
                         )
                     )
                 elif (
-                    datetime_based_cursor_component_definition
-                    and datetime_based_cursor_component_definition.get("type", "")
+                    incremental_sync_component_definition
+                    and incremental_sync_component_definition.get("type", "")
                     == DatetimeBasedCursorModel.__name__
                     and self._stream_supports_concurrent_partition_processing(
                         declarative_stream=declarative_stream
@@ -320,18 +321,16 @@ def _group_streams(
                     )
                     partition_router = declarative_stream.retriever.stream_slicer._partition_router
 
-                    cursor, connector_state_converter = (
-                        self._constructor.create_concurrent_cursor_from_perpartition_cursor(
+                    cursor = self._constructor.create_concurrent_cursor_from_perpartition_cursor(
                             state_manager=state_manager,
                             model_type=DatetimeBasedCursorModel,
-                            component_definition=datetime_based_cursor_component_definition,
+                            component_definition=incremental_sync_component_definition,
                             stream_name=declarative_stream.name,
                             stream_namespace=declarative_stream.namespace,
                             config=config or {},
                             stream_state=stream_state,
                             partition_router=partition_router,
                         )
-                    )
 
 
                     partition_generator = StreamSlicerPartitionGenerator(

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -1,46 +1,31 @@
 import copy
 
-from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
-
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
 import logging
 from collections import OrderedDict
-from typing import Any, Callable, Iterable, Mapping, MutableMapping, Optional, Union
+from typing import Any, Callable, Iterable, Mapping, MutableMapping, Optional
 
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 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.message import MessageRepository
 from airbyte_cdk.sources.streams.checkpoint.per_partition_key_serializer import (
     PerPartitionKeySerializer,
 )
-from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, CursorField, CursorValueType, GapType
+from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, CursorField
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.types import Record, StreamSlice, StreamState
-import functools
-from abc import ABC, abstractmethod
-from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional, Protocol, Tuple
-
-from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
-from airbyte_cdk.sources.message import MessageRepository
-from airbyte_cdk.sources.streams import NO_CURSOR_STATE_KEY
-from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
-from airbyte_cdk.sources.streams.concurrent.partitions.stream_slicer import StreamSlicer
-from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter import (
-    AbstractStreamStateConverter,
-)
-from airbyte_cdk.sources.types import StreamSlice
 
 logger = logging.getLogger("airbyte")
 
 
 class ConcurrentCursorFactory:
-    def __init__(self, create_function: Callable[[Mapping[str, Any]], DeclarativeCursor]):
+    def __init__(self, create_function: Callable[..., tuple[Cursor, ...]]):
         self._create_function = create_function
 
-    def create(self, stream_state: Mapping[str, Any]) -> DeclarativeCursor:
+    def create(self, stream_state: Mapping[str, Any]) -> Cursor:
         return self._create_function(stream_state=stream_state)[0]
 
 
@@ -78,52 +63,36 @@ def __init__(
         stream_state: Any,
         message_repository: MessageRepository,
         connector_state_manager: ConnectorStateManager,
-        connector_state_converter: AbstractStreamStateConverter,
         cursor_field: CursorField,
-        slice_boundary_fields: Optional[Tuple[str, str]],
-        start: Optional[CursorValueType],
-        end_provider: Callable[[], CursorValueType],
-        lookback_window: Optional[GapType] = None,
-        slice_range: Optional[GapType] = None,
-        cursor_granularity: Optional[GapType] = None,
     ) -> None:
         self._stream_name = stream_name
         self._stream_namespace = stream_namespace
         self._message_repository = message_repository
-        self._connector_state_converter = connector_state_converter
         self._connector_state_manager = connector_state_manager
         self._cursor_field = cursor_field
-        # To see some example where the slice boundaries might not be defined, check https://github.com/airbytehq/airbyte/blob/1ce84d6396e446e1ac2377362446e3fb94509461/airbyte-integrations/connectors/source-stripe/source_stripe/streams.py#L363-L379
-        self._slice_boundary_fields = slice_boundary_fields
-        self._start = start
-        self._end_provider = end_provider
-        self._conccurent_state = stream_state
-        # self.start, self._concurrent_state = self._get_concurrent_state(stream_state)
-        self._lookback_window = lookback_window
-        self._slice_range = slice_range
-        self._most_recent_cursor_value_per_partition: MutableMapping[Partition, Any] = {}
-        self._has_closed_at_least_one_slice = False
-        self._cursor_granularity = cursor_granularity
 
         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, Cursor] = OrderedDict()
         self._over_limit = 0
-        self._state = {}
         self._partition_serializer = PerPartitionKeySerializer()
 
         self._set_initial_state(stream_state)
 
+    @property
+    def cursor_field(self) -> CursorField:
+        return self._cursor_field
+
     @property
     def state(self) -> MutableMapping[str, Any]:
         states = []
         for partition_tuple, cursor in self._cursor_per_partition.items():
             cursor_state = cursor._connector_state_converter.convert_to_state_message(
                 cursor._cursor_field, cursor.state
             )
-            # print(cursor_state, cursor.state)
             if cursor_state:
                 states.append(
                     {
@@ -132,13 +101,8 @@ def state(self) -> MutableMapping[str, Any]:
                     }
                 )
         state: dict[str, Any] = {"states": states}
-        # print(state)
         return state
 
-    @property
-    def cursor_field(self) -> CursorField:
-        return self._cursor_field
-
     def close_partition(self, partition: Partition) -> None:
         self._cursor_per_partition[self._to_partition_key(partition._stream_slice.partition)].close_partition_without_emit(partition=partition)
 
@@ -253,8 +217,7 @@ def _set_initial_state(self, stream_state: StreamState) -> None:
         self._partition_router.set_initial_state(stream_state)
 
     def observe(self, record: Record) -> None:
-        print(f"ESTATE: {self._to_partition_key(record.partition._stream_slice.partition)}: {record.data[self.cursor_field.cursor_field_key]}")
-        self._cursor_per_partition[self._to_partition_key(record.partition._stream_slice.partition)].observe(record)
+        self._cursor_per_partition[self._to_partition_key(record.associated_slice.partition)].observe(record)
 
     def _to_partition_key(self, partition: Mapping[str, Any]) -> str:
         return self._partition_serializer.to_partition_key(partition)
@@ -267,9 +230,7 @@ def _create_cursor(self, cursor_state: Any) -> DeclarativeCursor:
         return cursor
 
     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(record)
 
     def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
         if not first.associated_slice or not second.associated_slice:
@@ -295,7 +256,7 @@ def _convert_record_to_cursor_record(record: Record) -> Record:
             else None,
         )
 
-    def _get_cursor(self, record: Record) -> DeclarativeCursor:
+    def _get_cursor(self, record: Record) -> Cursor:
         if not record.associated_slice:
             raise ValueError(
                 "Invalid state as stream slices that are emitted should refer to an existing cursor"

airbyte_cdk/sources/declarative/incremental/per_partition_cursor.py

@@ -244,15 +244,6 @@ def get_request_headers(
         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:
-                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(stream_slice.partition)] = cursor
             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={}),