airbyte-cdk 0.52.7__py3-none-any.whl → 0.52.8__py3-none-any.whl

airbyte_cdk/destinations/vector_db_based/config.py

@@ -85,6 +85,7 @@ class ProcessingConfigModel(BaseModel):
         ...,
         title="Chunk size",
         maximum=8191,
+        minimum=1,
         description="Size of chunks in tokens to store in vector store (make sure it is not too big for the context if your LLM)",
     )
     chunk_overlap: int = Field(

airbyte_cdk/sources/abstract_source.py

@@ -225,56 +225,18 @@ class AbstractSource(Source, ABC):
 
         if stream_state and "state" in dir(stream_instance):
             stream_instance.state = stream_state  # type: ignore # we check that state in the dir(stream_instance)
-            logger.info(f"Setting state of {stream_name} stream to {stream_state}")
-
-        slices = stream_instance.stream_slices(
-            cursor_field=configured_stream.cursor_field,
-            sync_mode=SyncMode.incremental,
-            stream_state=stream_state,
-        )
-        logger.debug(f"Processing stream slices for {stream_name} (sync_mode: incremental)", extra={"stream_slices": slices})
-
-        total_records_counter = 0
-        has_slices = False
-        for _slice in slices:
-            has_slices = True
-            if self._slice_logger.should_log_slice_message(logger):
-                yield self._slice_logger.create_slice_log_message(_slice)
-            records = stream_instance.read_records(
-                sync_mode=SyncMode.incremental,
-                stream_slice=_slice,
-                stream_state=stream_state,
-                cursor_field=configured_stream.cursor_field or None,
-            )
-            record_counter = 0
-            for message_counter, record_data_or_message in enumerate(records, start=1):
-                message = self._get_message(record_data_or_message, stream_instance)
-                yield from self._emit_queued_messages()
-                yield message
-                if message.type == MessageType.RECORD:
-                    record = message.record
-                    stream_state = stream_instance.get_updated_state(stream_state, record.data)
-                    checkpoint_interval = stream_instance.state_checkpoint_interval
-                    record_counter += 1
-                    if checkpoint_interval and record_counter % checkpoint_interval == 0:
-                        yield self._checkpoint_state(stream_instance, stream_state, state_manager)
-
-                    total_records_counter += 1
-                    # This functionality should ideally live outside of this method
-                    # but since state is managed inside this method, we keep track
-                    # of it here.
-                    if internal_config.is_limit_reached(total_records_counter):
-                        # Break from slice loop to save state and exit from _read_incremental function.
-                        break
-
-            yield self._checkpoint_state(stream_instance, stream_state, state_manager)
-            if internal_config.is_limit_reached(total_records_counter):
-                return
-
-        if not has_slices:
-            # Safety net to ensure we always emit at least one state message even if there are no slices
-            checkpoint = self._checkpoint_state(stream_instance, stream_state, state_manager)
-            yield checkpoint
+            logger.info(f"Setting state of {self.name} stream to {stream_state}")
+
+        for record_data_or_message in stream_instance.read_incremental(
+            configured_stream.cursor_field,
+            logger,
+            self._slice_logger,
+            stream_state,
+            state_manager,
+            self.per_stream_state_enabled,
+            internal_config,
+        ):
+            yield self._get_message(record_data_or_message, stream_instance)
 
     def _emit_queued_messages(self) -> Iterable[AirbyteMessage]:
         if self.message_repository:
@@ -297,17 +259,6 @@ class AbstractSource(Source, ABC):
                 if internal_config.is_limit_reached(total_records_counter):
                     return
 
-    def _checkpoint_state(self, stream: Stream, stream_state: Mapping[str, Any], state_manager: ConnectorStateManager) -> AirbyteMessage:
-        # First attempt to retrieve the current state using the stream's state property. We receive an AttributeError if the state
-        # property is not implemented by the stream instance and as a fallback, use the stream_state retrieved from the stream
-        # instance's deprecated get_updated_state() method.
-        try:
-            state_manager.update_state_for_stream(stream.name, stream.namespace, stream.state)  # type: ignore # we know the field might not exist...
-
-        except AttributeError:
-            state_manager.update_state_for_stream(stream.name, stream.namespace, stream_state)
-        return state_manager.create_state_message(stream.name, stream.namespace, send_per_stream_state=self.per_stream_state_enabled)
-
     @staticmethod
     def _apply_log_level_to_stream_logger(logger: logging.Logger, stream_instance: Stream) -> None:
         """

airbyte_cdk/sources/message/repository.py

@@ -75,12 +75,6 @@ class InMemoryMessageRepository(MessageRepository):
         self._log_level = log_level
 
     def emit_message(self, message: AirbyteMessage) -> None:
-        """
-        :param message: As of today, only AirbyteControlMessages are supported given that supporting other types of message will need more
-          work and therefore this work has been postponed
-        """
-        if message.type not in _SUPPORTED_MESSAGE_TYPES:
-            raise ValueError(f"As of today, only {_SUPPORTED_MESSAGE_TYPES} are supported as part of the InMemoryMessageRepository")
         self._message_queue.append(message)
 
     def log_message(self, level: Level, message_provider: Callable[[], LogMessage]) -> None:

airbyte_cdk/sources/source.py

@@ -6,7 +6,7 @@
 import logging
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from typing import Any, Generic, Iterable, List, Mapping, MutableMapping, TypeVar, Union
+from typing import Any, Dict, Generic, Iterable, List, Mapping, MutableMapping, Optional, TypeVar, Union
 
 from airbyte_cdk.connector import BaseConnector, DefaultConnectorMixin, TConfig
 from airbyte_cdk.models import AirbyteCatalog, AirbyteMessage, AirbyteStateMessage, AirbyteStateType, ConfiguredAirbyteCatalog
@@ -25,7 +25,7 @@ class BaseSource(BaseConnector[TConfig], ABC, Generic[TConfig, TState, TCatalog]
         ...
 
     @abstractmethod
-    def read(self, logger: logging.Logger, config: TConfig, catalog: TCatalog, state: TState = None) -> Iterable[AirbyteMessage]:
+    def read(self, logger: logging.Logger, config: TConfig, catalog: TCatalog, state: Optional[TState] = None) -> Iterable[AirbyteMessage]:
         """
         Returns a generator of the AirbyteMessages generated by reading the source with the given configuration, catalog, and state.
         """
@@ -43,8 +43,9 @@ class Source(
     BaseSource[Mapping[str, Any], Union[List[AirbyteStateMessage], MutableMapping[str, Any]], ConfiguredAirbyteCatalog],
     ABC,
 ):
-    # can be overridden to change an input state
-    def read_state(self, state_path: str) -> Union[List[AirbyteStateMessage], MutableMapping[str, Any]]:
+    # can be overridden to change an input state.
+    @classmethod
+    def read_state(cls, state_path: str) -> Union[List[AirbyteStateMessage], MutableMapping[str, Any]]:
         """
         Retrieves the input state of a sync by reading from the specified JSON file. Incoming state can be deserialized into either
         a JSON object for legacy state input or as a list of AirbyteStateMessages for the per-stream state format. Regardless of the
@@ -53,30 +54,30 @@ class Source(
         :return: The complete stream state based on the connector's previous sync
         """
         if state_path:
-            state_obj = self._read_json_file(state_path)
+            state_obj = BaseConnector._read_json_file(state_path)
             if not state_obj:
-                return self._emit_legacy_state_format({})
-            is_per_stream_state = isinstance(state_obj, List)
-            if is_per_stream_state:
+                return cls._emit_legacy_state_format({})
+            if isinstance(state_obj, List):
                 parsed_state_messages = []
-                for state in state_obj:
+                for state in state_obj:  # type: ignore  # `isinstance(state_obj, List)` ensures that this is a list
                     parsed_message = AirbyteStateMessage.parse_obj(state)
                     if not parsed_message.stream and not parsed_message.data and not parsed_message.global_:
                         raise ValueError("AirbyteStateMessage should contain either a stream, global, or state field")
                     parsed_state_messages.append(parsed_message)
                 return parsed_state_messages
             else:
-                return self._emit_legacy_state_format(state_obj)
-        return self._emit_legacy_state_format({})
+                return cls._emit_legacy_state_format(state_obj)  # type: ignore  # assuming it is a dict
+        return cls._emit_legacy_state_format({})
 
-    def _emit_legacy_state_format(self, state_obj) -> Union[List[AirbyteStateMessage], MutableMapping[str, Any]]:
+    @classmethod
+    def _emit_legacy_state_format(cls, state_obj: Dict[str, Any]) -> Union[List[AirbyteStateMessage], MutableMapping[str, Any]]:
         """
         Existing connectors that override read() might not be able to interpret the new state format. We temporarily
         send state in the old format for these connectors, but once all have been upgraded, this method can be removed,
         and we can then emit state in the list format.
         """
         # vars(self.__class__) checks if the current class directly overrides the read() function
-        if "read" in vars(self.__class__):
+        if "read" in vars(cls):
             return defaultdict(dict, state_obj)
         else:
             if state_obj:

airbyte_cdk/sources/streams/concurrent/adapters.py

@@ -6,10 +6,11 @@ import copy
 import json
 import logging
 from functools import lru_cache
-from typing import Any, Iterable, List, Mapping, Optional, Tuple, Union
+from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union
 
 from airbyte_cdk.models import AirbyteStream, SyncMode
 from airbyte_cdk.sources import AbstractSource, Source
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.availability_strategy import AvailabilityStrategy
@@ -20,12 +21,14 @@ from airbyte_cdk.sources.streams.concurrent.availability_strategy import (
     StreamAvailable,
     StreamUnavailable,
 )
+from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, NoopCursor
 from airbyte_cdk.sources.streams.concurrent.exceptions import ExceptionWithDisplayMessage
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 from airbyte_cdk.sources.streams.concurrent.thread_based_concurrent_stream import ThreadBasedConcurrentStream
 from airbyte_cdk.sources.streams.core import StreamData
+from airbyte_cdk.sources.utils.schema_helpers import InternalConfig
 from airbyte_cdk.sources.utils.slice_logger import SliceLogger
 from deprecated.classic import deprecated
 
@@ -44,7 +47,15 @@ class StreamFacade(Stream):
     """
 
     @classmethod
-    def create_from_stream(cls, stream: Stream, source: AbstractSource, logger: logging.Logger, max_workers: int) -> Stream:
+    def create_from_stream(
+        cls,
+        stream: Stream,
+        source: AbstractSource,
+        logger: logging.Logger,
+        max_workers: int,
+        state: Optional[MutableMapping[str, Any]],
+        cursor: Cursor,
+    ) -> Stream:
         """
         Create a ConcurrentStream from a Stream object.
         :param source: The source
@@ -63,9 +74,16 @@ class StreamFacade(Stream):
         message_repository = source.message_repository
         return StreamFacade(
             ThreadBasedConcurrentStream(
-                partition_generator=StreamPartitionGenerator(stream, message_repository),
+                partition_generator=StreamPartitionGenerator(
+                    stream,
+                    message_repository,
+                    SyncMode.full_refresh if isinstance(cursor, NoopCursor) else SyncMode.incremental,
+                    [cursor_field] if cursor_field is not None else None,
+                    state,
+                ),
                 max_workers=max_workers,
                 name=stream.name,
+                namespace=stream.namespace,
                 json_schema=stream.get_json_schema(),
                 availability_strategy=StreamAvailabilityStrategy(stream, source),
                 primary_key=pk,
@@ -73,9 +91,21 @@ class StreamFacade(Stream):
                 slice_logger=source._slice_logger,
                 message_repository=message_repository,
                 logger=logger,
-            )
+                cursor=cursor,
+            ),
+            stream,
+            cursor,
         )
 
+    @property
+    def state(self) -> MutableMapping[str, Any]:
+        raise NotImplementedError("This should not be called as part of the Concurrent CDK code. Please report the problem to Airbyte")
+
+    @state.setter
+    def state(self, value: Mapping[str, Any]) -> None:
+        if "state" in dir(self._legacy_stream):
+            self._legacy_stream.state = value  # type: ignore  # validating `state` is attribute of stream using `if` above
+
     @classmethod
     def _get_primary_key_from_stream(cls, stream_primary_key: Optional[Union[str, List[str], List[List[str]]]]) -> List[str]:
         if stream_primary_key is None:
@@ -102,11 +132,13 @@ class StreamFacade(Stream):
         else:
             return stream.cursor_field
 
-    def __init__(self, stream: AbstractStream):
+    def __init__(self, stream: AbstractStream, legacy_stream: Stream, cursor: Cursor):
         """
         :param stream: The underlying AbstractStream
         """
         self._abstract_stream = stream
+        self._legacy_stream = legacy_stream
+        self._cursor = cursor
 
     def read_full_refresh(
         self,
@@ -121,8 +153,19 @@ class StreamFacade(Stream):
         :param slice_logger: (ignored)
         :return: Iterable of StreamData
         """
-        for record in self._abstract_stream.read():
-            yield record.data
+        yield from self._read_records()
+
+    def read_incremental(
+        self,
+        cursor_field: Optional[List[str]],
+        logger: logging.Logger,
+        slice_logger: SliceLogger,
+        stream_state: MutableMapping[str, Any],
+        state_manager: ConnectorStateManager,
+        per_stream_state_enabled: bool,
+        internal_config: InternalConfig,
+    ) -> Iterable[StreamData]:
+        yield from self._read_records()
 
     def read_records(
         self,
@@ -131,12 +174,11 @@ class StreamFacade(Stream):
         stream_slice: Optional[Mapping[str, Any]] = None,
         stream_state: Optional[Mapping[str, Any]] = None,
     ) -> Iterable[StreamData]:
-        if sync_mode == SyncMode.full_refresh:
-            for record in self._abstract_stream.read():
-                yield record.data
-        else:
-            # Incremental reads are not supported
-            raise NotImplementedError
+        yield from self._read_records()
+
+    def _read_records(self) -> Iterable[StreamData]:
+        for record in self._abstract_stream.read():
+            yield record.data
 
     @property
     def name(self) -> str:
@@ -165,8 +207,7 @@ class StreamFacade(Stream):
 
     @property
     def supports_incremental(self) -> bool:
-        # Only full refresh is supported
-        return False
+        return self._legacy_stream.supports_incremental
 
     def check_availability(self, logger: logging.Logger, source: Optional["Source"] = None) -> Tuple[bool, Optional[str]]:
         """
@@ -210,7 +251,15 @@ class StreamPartition(Partition):
     In the long-run, it would be preferable to update the connectors, but we don't have the tooling or need to justify the effort at this time.
     """
 
-    def __init__(self, stream: Stream, _slice: Optional[Mapping[str, Any]], message_repository: MessageRepository):
+    def __init__(
+        self,
+        stream: Stream,
+        _slice: Optional[Mapping[str, Any]],
+        message_repository: MessageRepository,
+        sync_mode: SyncMode,
+        cursor_field: Optional[List[str]],
+        state: Optional[MutableMapping[str, Any]],
+    ):
         """
         :param stream: The stream to delegate to
         :param _slice: The partition's stream_slice
@@ -219,6 +268,9 @@ class StreamPartition(Partition):
         self._stream = stream
         self._slice = _slice
         self._message_repository = message_repository
+        self._sync_mode = sync_mode
+        self._cursor_field = cursor_field
+        self._state = state
 
     def read(self) -> Iterable[Record]:
         """
@@ -227,7 +279,18 @@ class StreamPartition(Partition):
         Otherwise, the message will be emitted on the message repository.
         """
         try:
-            for record_data in self._stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=copy.deepcopy(self._slice)):
+            # using `stream_state=self._state` have a very different behavior than the current one as today the state is updated slice
+            #  by slice incrementally. We don't have this guarantee with Concurrent CDK. For HttpStream, `stream_state` is passed to:
+            #  * fetch_next_page
+            #  * parse_response
+            #  Both are not used for Stripe so we should be good for the first iteration of Concurrent CDK. However, Stripe still do
+            #  `if not stream_state` to know if it calls the Event stream or not
+            for record_data in self._stream.read_records(
+                cursor_field=self._cursor_field,
+                sync_mode=SyncMode.full_refresh,
+                stream_slice=copy.deepcopy(self._slice),
+                stream_state=self._state,
+            ):
                 if isinstance(record_data, Mapping):
                     data_to_return = dict(record_data)
                     self._stream.transformer.transform(data_to_return, self._stream.get_json_schema())
@@ -264,17 +327,27 @@ class StreamPartitionGenerator(PartitionGenerator):
     In the long-run, it would be preferable to update the connectors, but we don't have the tooling or need to justify the effort at this time.
     """
 
-    def __init__(self, stream: Stream, message_repository: MessageRepository):
+    def __init__(
+        self,
+        stream: Stream,
+        message_repository: MessageRepository,
+        sync_mode: SyncMode,
+        cursor_field: Optional[List[str]],
+        state: Optional[MutableMapping[str, Any]],
+    ):
         """
         :param stream: The stream to delegate to
         :param message_repository: The message repository to use to emit non-record messages
         """
         self.message_repository = message_repository
         self._stream = stream
+        self._sync_mode = sync_mode
+        self._cursor_field = cursor_field
+        self._state = state
 
-    def generate(self, sync_mode: SyncMode) -> Iterable[Partition]:
-        for s in self._stream.stream_slices(sync_mode=sync_mode):
-            yield StreamPartition(self._stream, copy.deepcopy(s), self.message_repository)
+    def generate(self) -> Iterable[Partition]:
+        for s in self._stream.stream_slices(sync_mode=self._sync_mode, cursor_field=self._cursor_field, stream_state=self._state):
+            yield StreamPartition(self._stream, copy.deepcopy(s), self.message_repository, self._sync_mode, self._cursor_field, self._state)
 
 
 @deprecated("This class is experimental. Use at your own risk.")

airbyte_cdk/sources/streams/concurrent/cursor.py

@@ -0,0 +1,148 @@
+import functools
+from abc import ABC, abstractmethod
+from typing import Any, List, Mapping, Optional, Protocol, Tuple
+
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
+from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
+from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+
+
+def _extract_value(mapping: Mapping[str, Any], path: List[str]) -> Any:
+    return functools.reduce(lambda a, b: a[b], path, mapping)
+
+
+class Comparable(Protocol):
+    """Protocol for annotating comparable types."""
+
+    @abstractmethod
+    def __lt__(self: "Comparable", other: "Comparable") -> bool:
+        pass
+
+
+class CursorField:
+    def __init__(self, cursor_field_key: str) -> None:
+        self._cursor_field_key = cursor_field_key
+
+    def extract_value(self, record: Record) -> Comparable:
+        cursor_value = record.data.get(self._cursor_field_key)
+        if cursor_value is None:
+            raise ValueError(f"Could not find cursor field {self._cursor_field_key} in record")
+        return cursor_value  # type: ignore  # we assume that the value the path points at is a comparable
+
+
+class Cursor(ABC):
+    @abstractmethod
+    def observe(self, record: Record) -> None:
+        """
+        Indicate to the cursor that the record has been emitted
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def close_partition(self, partition: Partition) -> None:
+        """
+        Indicate to the cursor that the partition has been successfully processed
+        """
+        raise NotImplementedError()
+
+
+class NoopCursor(Cursor):
+    def observe(self, record: Record) -> None:
+        pass
+
+    def close_partition(self, partition: Partition) -> None:
+        pass
+
+
+class ConcurrentCursor(Cursor):
+    _START_BOUNDARY = 0
+    _END_BOUNDARY = 1
+
+    def __init__(
+        self,
+        stream_name: str,
+        stream_namespace: Optional[str],
+        stream_state: Any,
+        message_repository: MessageRepository,
+        connector_state_manager: ConnectorStateManager,
+        cursor_field: CursorField,
+        slice_boundary_fields: Optional[Tuple[str, str]],
+    ) -> None:
+        self._stream_name = stream_name
+        self._stream_namespace = stream_namespace
+        self._message_repository = message_repository
+        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 if slice_boundary_fields else tuple()
+        self._most_recent_record: Optional[Record] = None
+        self._has_closed_at_least_one_slice = False
+
+        # TODO to migrate state. The migration should probably be outside of this class. Impact of not having this:
+        #  * Given a sync that emits no records, the emitted state message will be empty
+        self._state = {  # type: ignore  # waiting for when we implement the state migration
+            "slices": []  # empty for now but should look like `{start: 1, end: 10, parent_id: "id1"}`
+        }
+
+    def observe(self, record: Record) -> None:
+        if self._slice_boundary_fields:
+            # Given that slicing is done using the cursor field, we don't need to observe the record as we assume slices will describe what
+            # has been emitted. Assuming there is a chance that records might not be yet populated for the most recent slice, use a lookback
+            # window
+            return
+
+        if not self._most_recent_record or self._extract_cursor_value(self._most_recent_record) < self._extract_cursor_value(record):
+            self._most_recent_record = record
+
+    def _extract_cursor_value(self, record: Record) -> Comparable:
+        return self._cursor_field.extract_value(record)
+
+    def close_partition(self, partition: Partition) -> None:
+        slice_count_before = len(self._state["slices"])
+        self._add_slice_to_state(partition)
+        if slice_count_before < len(self._state["slices"]):
+            self._merge_partitions()
+            self._emit_state_message()
+        self._has_closed_at_least_one_slice = True
+
+    def _add_slice_to_state(self, partition: Partition) -> None:
+        if self._slice_boundary_fields:
+            self._state["slices"].append(
+                {
+                    "start": self._extract_from_slice(partition, self._slice_boundary_fields[self._START_BOUNDARY]),
+                    "end": self._extract_from_slice(partition, self._slice_boundary_fields[self._END_BOUNDARY]),
+                }
+            )
+        elif self._most_recent_record:
+            if self._has_closed_at_least_one_slice:
+                raise ValueError(
+                    "Given that slice_boundary_fields is not defined and that per-partition state is not supported, only one slice is "
+                    "expected."
+                )
+
+            self._state["slices"].append(
+                {
+                    "start": 0,  # FIXME this only works with int datetime
+                    "end": self._extract_cursor_value(self._most_recent_record),
+                }
+            )
+
+    def _emit_state_message(self) -> None:
+        self._connector_state_manager.update_state_for_stream(self._stream_name, self._stream_namespace, self._state)
+        state_message = self._connector_state_manager.create_state_message(
+            self._stream_name, self._stream_namespace, send_per_stream_state=True
+        )
+        self._message_repository.emit_message(state_message)
+
+    def _merge_partitions(self) -> None:
+        pass  # TODO eventually
+
+    def _extract_from_slice(self, partition: Partition, key: str) -> Comparable:
+        try:
+            _slice = partition.to_slice()
+            if not _slice:
+                raise KeyError(f"Could not find key `{key}` in empty slice")
+            return _slice[key]  # type: ignore  # we expect the devs to specify a key that would return a Comparable
+        except KeyError as exception:
+            raise KeyError(f"Partition is expected to have key `{key}` but could not be found") from exception

airbyte_cdk/sources/streams/concurrent/partition_enqueuer.py

@@ -4,7 +4,6 @@
 
 from queue import Queue
 
-from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
 from airbyte_cdk.sources.streams.concurrent.partitions.types import PARTITIONS_GENERATED_SENTINEL, QueueItem
 
@@ -22,7 +21,7 @@ class PartitionEnqueuer:
         self._queue = queue
         self._sentinel = sentinel
 
-    def generate_partitions(self, partition_generator: PartitionGenerator, sync_mode: SyncMode) -> None:
+    def generate_partitions(self, partition_generator: PartitionGenerator) -> None:
         """
         Generate partitions from a partition generator and put them in a queue.
         When all the partitions are added to the queue, a sentinel is added to the queue to indicate that all the partitions have been generated.
@@ -35,7 +34,7 @@ class PartitionEnqueuer:
         :return:
         """
         try:
-            for partition in partition_generator.generate(sync_mode=sync_mode):
+            for partition in partition_generator.generate():
                 self._queue.put(partition)
             self._queue.put(self._sentinel)
         except Exception as e:

airbyte_cdk/sources/streams/concurrent/partitions/partition.py

@@ -25,6 +25,9 @@ class Partition(ABC):
     def to_slice(self) -> Optional[Mapping[str, Any]]:
         """
         Converts the partition to a slice that can be serialized and deserialized.
+
+        Note: it would have been interesting to have a type of `Mapping[str, Comparable]` to simplify typing but some slices can have nested
+         values ([example](https://github.com/airbytehq/airbyte/blob/1ce84d6396e446e1ac2377362446e3fb94509461/airbyte-integrations/connectors/source-stripe/source_stripe/streams.py#L584-L596))
         :return: A mapping representing a slice
         """
         pass

airbyte_cdk/sources/streams/concurrent/partitions/partition_generator.py

@@ -5,16 +5,14 @@
 from abc import ABC, abstractmethod
 from typing import Iterable
 
-from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 
 
 class PartitionGenerator(ABC):
     @abstractmethod
-    def generate(self, sync_mode: SyncMode) -> Iterable[Partition]:
+    def generate(self) -> Iterable[Partition]:
         """
         Generates partitions for a given sync mode.
-        :param sync_mode: SyncMode
         :return: An iterable of partitions
         """
         pass

airbyte_cdk/sources/streams/concurrent/thread_based_concurrent_stream.py

@@ -14,6 +14,7 @@ from airbyte_cdk.models import AirbyteStream, SyncMode
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import AbstractAvailabilityStrategy, StreamAvailability
+from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, NoopCursor
 from airbyte_cdk.sources.streams.concurrent.partition_enqueuer import PartitionEnqueuer
 from airbyte_cdk.sources.streams.concurrent.partition_reader import PartitionReader
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
@@ -44,6 +45,7 @@ class ThreadBasedConcurrentStream(AbstractStream):
         timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
         max_concurrent_tasks: int = DEFAULT_MAX_QUEUE_SIZE,
         sleep_time: float = DEFAULT_SLEEP_TIME,
+        cursor: Cursor = NoopCursor(),
         namespace: Optional[str] = None,
     ):
         self._stream_partition_generator = partition_generator
@@ -60,6 +62,7 @@ class ThreadBasedConcurrentStream(AbstractStream):
         self._timeout_seconds = timeout_seconds
         self._max_concurrent_tasks = max_concurrent_tasks
         self._sleep_time = sleep_time
+        self._cursor = cursor
         self._namespace = namespace
 
     def read(self) -> Iterable[Record]:
@@ -80,14 +83,13 @@ class ThreadBasedConcurrentStream(AbstractStream):
           - If the next work item is a PartitionCompleteSentinel, a partition is done processing.
             - Update the value in partitions_to_done to True so we know the partition is completed.
         """
-        self._logger.debug(f"Processing stream slices for {self.name} (sync_mode: full_refresh)")
+        self._logger.debug(f"Processing stream slices for {self.name}")
         futures: List[Future[Any]] = []
         queue: Queue[QueueItem] = Queue()
         partition_generator = PartitionEnqueuer(queue, PARTITIONS_GENERATED_SENTINEL)
         partition_reader = PartitionReader(queue)
 
-        # Submit partition generation tasks
-        self._submit_task(futures, partition_generator.generate_partitions, self._stream_partition_generator, SyncMode.full_refresh)
+        self._submit_task(futures, partition_generator.generate_partitions, self._stream_partition_generator)
 
         # True -> partition is done
         # False -> partition is not done
@@ -109,9 +111,11 @@ class ThreadBasedConcurrentStream(AbstractStream):
                         f"Received sentinel for partition {record_or_partition_or_exception.partition} that was not in partitions. This is indicative of a bug in the CDK. Please contact support.partitions:\n{partitions_to_done}"
                     )
                 partitions_to_done[record_or_partition_or_exception.partition] = True
+                self._cursor.close_partition(record_or_partition_or_exception.partition)
             elif isinstance(record_or_partition_or_exception, Record):
                 # Emit records
                 yield record_or_partition_or_exception
+                self._cursor.observe(record_or_partition_or_exception)
             elif isinstance(record_or_partition_or_exception, Partition):
                 # A new partition was generated and must be processed
                 partitions_to_done[record_or_partition_or_exception] = False