airbyte-cdk 0.72.0__py3-none-any.whl → 6.17.1.dev0__py3-none-any.whl

airbyte_cdk/sources/streams/concurrent/cursor.py

@@ -1,28 +1,67 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
+
 import functools
+import logging
 from abc import ABC, abstractmethod
-from datetime import datetime
-from typing import Any, List, Mapping, MutableMapping, Optional, Protocol, Tuple
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Protocol,
+    Tuple,
+    Union,
+)
 
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.message import MessageRepository
-from airbyte_cdk.sources.streams import FULL_REFRESH_SENTINEL_STATE_KEY
+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.state_converters.abstract_stream_state_converter import AbstractStreamStateConverter
+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 Record, StreamSlice
+
+LOGGER = logging.getLogger("airbyte")
 
 
 def _extract_value(mapping: Mapping[str, Any], path: List[str]) -> Any:
     return functools.reduce(lambda a, b: a[b], path, mapping)
 
 
-class Comparable(Protocol):
+class GapType(Protocol):
+    """
+    This is the representation of gaps between two cursor values. Examples:
+    * if cursor values are datetimes, GapType is timedelta
+    * if cursor values are integer, GapType will also be integer
+    """
+
+    pass
+
+
+class CursorValueType(Protocol):
     """Protocol for annotating comparable types."""
 
     @abstractmethod
-    def __lt__(self: "Comparable", other: "Comparable") -> bool:
+    def __lt__(self: "CursorValueType", other: "CursorValueType") -> bool:
+        pass
+
+    @abstractmethod
+    def __ge__(self: "CursorValueType", other: "CursorValueType") -> bool:
+        pass
+
+    @abstractmethod
+    def __add__(self: "CursorValueType", other: GapType) -> "CursorValueType":
+        pass
+
+    @abstractmethod
+    def __sub__(self: "CursorValueType", other: GapType) -> "CursorValueType":
         pass
 
 
@@ -30,18 +69,17 @@ class CursorField:
     def __init__(self, cursor_field_key: str) -> None:
         self.cursor_field_key = cursor_field_key
 
-    def extract_value(self, record: Record) -> Comparable:
+    def extract_value(self, record: Record) -> CursorValueType:
         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):
+class Cursor(StreamSlicer, ABC):
     @property
     @abstractmethod
-    def state(self) -> MutableMapping[str, Any]:
-        ...
+    def state(self) -> MutableMapping[str, Any]: ...
 
     @abstractmethod
     def observe(self, record: Record) -> None:
@@ -65,6 +103,13 @@ class Cursor(ABC):
         """
         raise NotImplementedError()
 
+    def stream_slices(self) -> Iterable[StreamSlice]:
+        """
+        Default placeholder implementation of generate_slices.
+        Subclasses can override this method to provide actual behavior.
+        """
+        yield StreamSlice(partition={}, cursor_slice={})
+
 
 class FinalStateCursor(Cursor):
     """Cursor that is used to guarantee at least one state message is emitted for a concurrent stream."""
@@ -81,12 +126,12 @@ class FinalStateCursor(Cursor):
         # Normally the connector state manager operates at the source-level. However, we only need it to write the sentinel
         # state message rather than manage overall source state. This is also only temporary as we move to the resumable
         # full refresh world where every stream uses a FileBasedConcurrentCursor with incremental state.
-        self._connector_state_manager = ConnectorStateManager(stream_instance_map={})
+        self._connector_state_manager = ConnectorStateManager()
         self._has_closed_at_least_one_slice = False
 
     @property
     def state(self) -> MutableMapping[str, Any]:
-        return {FULL_REFRESH_SENTINEL_STATE_KEY: True}
+        return {NO_CURSOR_STATE_KEY: True}
 
     def observe(self, record: Record) -> None:
         pass
@@ -99,8 +144,12 @@ class FinalStateCursor(Cursor):
         Used primarily for full refresh syncs that do not have a valid cursor value to emit at the end of a sync
         """
 
-        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)
+        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
+        )
         self._message_repository.emit_message(state_message)
 
 
@@ -118,7 +167,11 @@ class ConcurrentCursor(Cursor):
         connector_state_converter: AbstractStreamStateConverter,
         cursor_field: CursorField,
         slice_boundary_fields: Optional[Tuple[str, str]],
-        start: Optional[Any],
+        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
@@ -127,30 +180,62 @@ class ConcurrentCursor(Cursor):
         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._slice_boundary_fields = slice_boundary_fields
         self._start = start
-        self._most_recent_record: Optional[Record] = None
-        self._has_closed_at_least_one_slice = False
+        self._end_provider = end_provider
         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[
+            Union[StreamSlice, Mapping[str, Any], None], Any
+        ] = {}
+        self._has_closed_at_least_one_slice = False
+        self._cursor_granularity = cursor_granularity
+        # Flag to track if the logger has been triggered (per stream)
+        self._should_be_synced_logger_triggered = False
 
     @property
     def state(self) -> MutableMapping[str, Any]:
         return self._concurrent_state
 
-    def _get_concurrent_state(self, state: MutableMapping[str, Any]) -> Tuple[datetime, MutableMapping[str, Any]]:
+    @property
+    def cursor_field(self) -> CursorField:
+        return self._cursor_field
+
+    @property
+    def _slice_boundary_fields_wrapper(self) -> Tuple[str, str]:
+        return (
+            self._slice_boundary_fields
+            if self._slice_boundary_fields
+            else (
+                self._connector_state_converter.START_KEY,
+                self._connector_state_converter.END_KEY,
+            )
+        )
+
+    def _get_concurrent_state(
+        self, state: MutableMapping[str, Any]
+    ) -> Tuple[CursorValueType, MutableMapping[str, Any]]:
         if self._connector_state_converter.is_state_message_compatible(state):
-            return self._start or self._connector_state_converter.zero_value, self._connector_state_converter.deserialize(state)
-        return self._connector_state_converter.convert_from_sequential_state(self._cursor_field, state, self._start)
+            return (
+                self._start or self._connector_state_converter.zero_value,
+                self._connector_state_converter.deserialize(state),
+            )
+        return self._connector_state_converter.convert_from_sequential_state(
+            self._cursor_field, state, self._start
+        )
 
     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
+        most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(
+            record.associated_slice
+        )
+        try:
+            cursor_value = self._extract_cursor_value(record)
 
-        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
+            if most_recent_cursor_value is None or most_recent_cursor_value < cursor_value:
+                self._most_recent_cursor_value_per_partition[record.associated_slice] = cursor_value
+        except ValueError:
+            self._log_for_record_without_cursor_value()
 
     def _extract_cursor_value(self, record: Record) -> Any:
         return self._connector_state_converter.parse_value(self._cursor_field.extract_value(record))
@@ -158,12 +243,18 @@ class ConcurrentCursor(Cursor):
     def close_partition(self, partition: Partition) -> None:
         slice_count_before = len(self.state.get("slices", []))
         self._add_slice_to_state(partition)
-        if slice_count_before < len(self.state["slices"]):  # only emit if at least one slice has been processed
+        if slice_count_before < len(
+            self.state["slices"]
+        ):  # only emit if at least one slice has been processed
             self._merge_partitions()
             self._emit_state_message()
         self._has_closed_at_least_one_slice = True
 
     def _add_slice_to_state(self, partition: Partition) -> None:
+        most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(
+            partition.to_slice()
+        )
+
         if self._slice_boundary_fields:
             if "slices" not in self.state:
                 raise RuntimeError(
@@ -171,11 +262,16 @@ class ConcurrentCursor(Cursor):
                 )
             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]),
+                    self._connector_state_converter.START_KEY: self._extract_from_slice(
+                        partition, self._slice_boundary_fields[self._START_BOUNDARY]
+                    ),
+                    self._connector_state_converter.END_KEY: self._extract_from_slice(
+                        partition, self._slice_boundary_fields[self._END_BOUNDARY]
+                    ),
+                    self._connector_state_converter.MOST_RECENT_RECORD_KEY: most_recent_cursor_value,
                 }
             )
-        elif self._most_recent_record:
+        elif most_recent_cursor_value:
             if self._has_closed_at_least_one_slice:
                 # If we track state value using records cursor field, we can only do that if there is one partition. This is because we save
                 # the state every time we close a partition. We assume that if there are multiple slices, they need to be providing
@@ -195,7 +291,8 @@ class ConcurrentCursor(Cursor):
             self.state["slices"].append(
                 {
                     self._connector_state_converter.START_KEY: self.start,
-                    self._connector_state_converter.END_KEY: self._extract_cursor_value(self._most_recent_record),
+                    self._connector_state_converter.END_KEY: most_recent_cursor_value,
+                    self._connector_state_converter.MOST_RECENT_RECORD_KEY: most_recent_cursor_value,
                 }
             )
 
@@ -203,25 +300,28 @@ class ConcurrentCursor(Cursor):
         self._connector_state_manager.update_state_for_stream(
             self._stream_name,
             self._stream_namespace,
-            self._connector_state_converter.convert_to_sequential_state(self._cursor_field, self.state),
+            self._connector_state_converter.convert_to_state_message(
+                self._cursor_field, self.state
+            ),
+        )
+        state_message = self._connector_state_manager.create_state_message(
+            self._stream_name, self._stream_namespace
         )
-        # TODO: if we migrate stored state to the concurrent state format
-        #  (aka stop calling self._connector_state_converter.convert_to_sequential_state`), we'll need to cast datetimes to string or
-        #  int before emitting state
-        state_message = self._connector_state_manager.create_state_message(self._stream_name, self._stream_namespace)
         self._message_repository.emit_message(state_message)
 
     def _merge_partitions(self) -> None:
         self.state["slices"] = self._connector_state_converter.merge_intervals(self.state["slices"])
 
-    def _extract_from_slice(self, partition: Partition, key: str) -> Comparable:
+    def _extract_from_slice(self, partition: Partition, key: str) -> CursorValueType:
         try:
             _slice = partition.to_slice()
             if not _slice:
                 raise KeyError(f"Could not find key `{key}` in empty slice")
-            return self._connector_state_converter.parse_value(_slice[key])  # type: ignore  # we expect the devs to specify a key that would return a Comparable
+            return self._connector_state_converter.parse_value(_slice[key])  # type: ignore  # we expect the devs to specify a key that would return a CursorValueType
         except KeyError as exception:
-            raise KeyError(f"Partition is expected to have key `{key}` but could not be found") from exception
+            raise KeyError(
+                f"Partition is expected to have key `{key}` but could not be found"
+            ) from exception
 
     def ensure_at_least_one_state_emitted(self) -> None:
         """
@@ -229,3 +329,159 @@ class ConcurrentCursor(Cursor):
         called.
         """
         self._emit_state_message()
+
+    def stream_slices(self) -> Iterable[StreamSlice]:
+        """
+        Generating slices based on a few parameters:
+        * lookback_window: Buffer to remove from END_KEY of the highest slice
+        * slice_range: Max difference between two slices. If the difference between two slices is greater, multiple slices will be created
+        * start: `_split_per_slice_range` will clip any value to `self._start which means that:
+          * if upper is less than self._start, no slices will be generated
+          * if lower is less than self._start, self._start will be used as the lower boundary (lookback_window will not be considered in that case)
+
+        Note that the slices will overlap at their boundaries. We therefore expect to have at least the lower or the upper boundary to be
+        inclusive in the API that is queried.
+        """
+        self._merge_partitions()
+
+        if self._start is not None and self._is_start_before_first_slice():
+            yield from self._split_per_slice_range(
+                self._start,
+                self.state["slices"][0][self._connector_state_converter.START_KEY],
+                False,
+            )
+
+        if len(self.state["slices"]) == 1:
+            yield from self._split_per_slice_range(
+                self._calculate_lower_boundary_of_last_slice(
+                    self.state["slices"][0][self._connector_state_converter.END_KEY]
+                ),
+                self._end_provider(),
+                True,
+            )
+        elif len(self.state["slices"]) > 1:
+            for i in range(len(self.state["slices"]) - 1):
+                if self._cursor_granularity:
+                    yield from self._split_per_slice_range(
+                        self.state["slices"][i][self._connector_state_converter.END_KEY]
+                        + self._cursor_granularity,
+                        self.state["slices"][i + 1][self._connector_state_converter.START_KEY],
+                        False,
+                    )
+                else:
+                    yield from self._split_per_slice_range(
+                        self.state["slices"][i][self._connector_state_converter.END_KEY],
+                        self.state["slices"][i + 1][self._connector_state_converter.START_KEY],
+                        False,
+                    )
+            yield from self._split_per_slice_range(
+                self._calculate_lower_boundary_of_last_slice(
+                    self.state["slices"][-1][self._connector_state_converter.END_KEY]
+                ),
+                self._end_provider(),
+                True,
+            )
+        else:
+            raise ValueError("Expected at least one slice")
+
+    def _is_start_before_first_slice(self) -> bool:
+        return (
+            self._start is not None
+            and self._start < self.state["slices"][0][self._connector_state_converter.START_KEY]
+        )
+
+    def _calculate_lower_boundary_of_last_slice(
+        self, lower_boundary: CursorValueType
+    ) -> CursorValueType:
+        if self._lookback_window:
+            return lower_boundary - self._lookback_window
+        return lower_boundary
+
+    def _split_per_slice_range(
+        self, lower: CursorValueType, upper: CursorValueType, upper_is_end: bool
+    ) -> Iterable[StreamSlice]:
+        if lower >= upper:
+            return
+
+        if self._start and upper < self._start:
+            return
+
+        lower = max(lower, self._start) if self._start else lower
+        if not self._slice_range or self._evaluate_upper_safely(lower, self._slice_range) >= upper:
+            start_value, end_value = (
+                (lower, upper - self._cursor_granularity)
+                if self._cursor_granularity and not upper_is_end
+                else (lower, upper)
+            )
+            yield StreamSlice(
+                partition={},
+                cursor_slice={
+                    self._slice_boundary_fields_wrapper[
+                        self._START_BOUNDARY
+                    ]: self._connector_state_converter.output_format(start_value),
+                    self._slice_boundary_fields_wrapper[
+                        self._END_BOUNDARY
+                    ]: self._connector_state_converter.output_format(end_value),
+                },
+            )
+        else:
+            stop_processing = False
+            current_lower_boundary = lower
+            while not stop_processing:
+                current_upper_boundary = min(
+                    self._evaluate_upper_safely(current_lower_boundary, self._slice_range), upper
+                )
+                has_reached_upper_boundary = current_upper_boundary >= upper
+
+                start_value, end_value = (
+                    (current_lower_boundary, current_upper_boundary - self._cursor_granularity)
+                    if self._cursor_granularity
+                    and (not upper_is_end or not has_reached_upper_boundary)
+                    else (current_lower_boundary, current_upper_boundary)
+                )
+                yield StreamSlice(
+                    partition={},
+                    cursor_slice={
+                        self._slice_boundary_fields_wrapper[
+                            self._START_BOUNDARY
+                        ]: self._connector_state_converter.output_format(start_value),
+                        self._slice_boundary_fields_wrapper[
+                            self._END_BOUNDARY
+                        ]: self._connector_state_converter.output_format(end_value),
+                    },
+                )
+                current_lower_boundary = current_upper_boundary
+                if current_upper_boundary >= upper:
+                    stop_processing = True
+
+    def _evaluate_upper_safely(self, lower: CursorValueType, step: GapType) -> CursorValueType:
+        """
+        Given that we set the default step at datetime.timedelta.max, we will generate an OverflowError when evaluating the next start_date
+        This method assumes that users would never enter a step that would generate an overflow. Given that would be the case, the code
+        would have broken anyway.
+        """
+        try:
+            return lower + step
+        except OverflowError:
+            return self._end_provider()
+
+    def should_be_synced(self, record: Record) -> bool:
+        """
+        Determines if a record should be synced based on its cursor value.
+        :param record: The record to evaluate
+
+        :return: True if the record's cursor value falls within the sync boundaries
+        """
+        try:
+            record_cursor_value: CursorValueType = self._extract_cursor_value(record)
+        except ValueError:
+            self._log_for_record_without_cursor_value()
+            return True
+        return self.start <= record_cursor_value <= self._end_provider()
+
+    def _log_for_record_without_cursor_value(self) -> None:
+        if not self._should_be_synced_logger_triggered:
+            LOGGER.warning(
+                f"Could not find cursor field `{self.cursor_field.cursor_field_key}` in record for stream {self._stream_name}. The incremental sync will assume it needs to be synced"
+            )
+            self._should_be_synced_logger_triggered = True

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -8,7 +8,10 @@ from typing import Any, Iterable, List, Mapping, Optional
 
 from airbyte_cdk.models import AirbyteStream, SyncMode
 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.availability_strategy import (
+    AbstractAvailabilityStrategy,
+    StreamAvailability,
+)
 from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
@@ -60,19 +63,25 @@ class DefaultStream(AbstractStream):
         return self._json_schema
 
     def as_airbyte_stream(self) -> AirbyteStream:
-        stream = AirbyteStream(name=self.name, json_schema=dict(self._json_schema), supported_sync_modes=[SyncMode.full_refresh])
+        stream = AirbyteStream(
+            name=self.name,
+            json_schema=dict(self._json_schema),
+            supported_sync_modes=[SyncMode.full_refresh],
+            is_resumable=False,
+        )
 
         if self._namespace:
             stream.namespace = self._namespace
 
         if self._cursor_field:
             stream.source_defined_cursor = True
+            stream.is_resumable = True
             stream.supported_sync_modes.append(SyncMode.incremental)
             stream.default_cursor_field = [self._cursor_field]
 
         keys = self._primary_key
         if keys and len(keys) > 0:
-            stream.source_defined_primary_key = [keys]
+            stream.source_defined_primary_key = [[key] for key in keys]
 
         return stream
 

airbyte_cdk/sources/streams/concurrent/exceptions.py

@@ -13,3 +13,6 @@ class ExceptionWithDisplayMessage(Exception):
     def __init__(self, display_message: str, **kwargs: Any):
         super().__init__(**kwargs)
         self.display_message = display_message
+
+    def __str__(self) -> str:
+        return f'ExceptionWithDisplayMessage: "{self.display_message}"'

airbyte_cdk/sources/streams/concurrent/helpers.py

@@ -5,14 +5,23 @@ from typing import List, Optional, Union
 from airbyte_cdk.sources.streams import Stream
 
 
-def get_primary_key_from_stream(stream_primary_key: Optional[Union[str, List[str], List[List[str]]]]) -> List[str]:
+def get_primary_key_from_stream(
+    stream_primary_key: Optional[Union[str, List[str], List[List[str]]]],
+) -> List[str]:
     if stream_primary_key is None:
         return []
     elif isinstance(stream_primary_key, str):
         return [stream_primary_key]
     elif isinstance(stream_primary_key, list):
-        if len(stream_primary_key) > 0 and all(isinstance(k, str) for k in stream_primary_key):
+        are_all_elements_str = all(isinstance(k, str) for k in stream_primary_key)
+        are_all_elements_list_of_size_one = all(
+            isinstance(k, list) and len(k) == 1 for k in stream_primary_key
+        )
+
+        if are_all_elements_str:
             return stream_primary_key  # type: ignore # We verified all items in the list are strings
+        elif are_all_elements_list_of_size_one:
+            return list(map(lambda x: x[0], stream_primary_key))
         else:
             raise ValueError(f"Nested primary keys are not supported. Found {stream_primary_key}")
     else:
@@ -22,7 +31,9 @@ def get_primary_key_from_stream(stream_primary_key: Optional[Union[str, List[str
 def get_cursor_field_from_stream(stream: Stream) -> Optional[str]:
     if isinstance(stream.cursor_field, list):
         if len(stream.cursor_field) > 1:
-            raise ValueError(f"Nested cursor fields are not supported. Got {stream.cursor_field} for {stream.name}")
+            raise ValueError(
+                f"Nested cursor fields are not supported. Got {stream.cursor_field} for {stream.name}"
+            )
         elif len(stream.cursor_field) == 0:
             return None
         else:

airbyte_cdk/sources/streams/concurrent/partition_enqueuer.py

@@ -4,7 +4,10 @@
 import time
 from queue import Queue
 
-from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import PartitionGenerationCompletedSentinel
+from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import (
+    PartitionGenerationCompletedSentinel,
+)
+from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.concurrent.partitions.types import QueueItem
@@ -15,7 +18,12 @@ class PartitionEnqueuer:
     Generates partitions from a partition generator and puts them in a queue.
     """
 
-    def __init__(self, queue: Queue[QueueItem], thread_pool_manager: ThreadPoolManager, sleep_time_in_seconds: float = 0.1) -> None:
+    def __init__(
+        self,
+        queue: Queue[QueueItem],
+        thread_pool_manager: ThreadPoolManager,
+        sleep_time_in_seconds: float = 0.1,
+    ) -> None:
         """
         :param queue:  The queue to put the partitions in.
         :param throttler: The throttler to use to throttle the partition generation.
@@ -52,4 +60,5 @@ class PartitionEnqueuer:
                 self._queue.put(partition)
             self._queue.put(PartitionGenerationCompletedSentinel(stream))
         except Exception as e:
-            self._queue.put(e)
+            self._queue.put(StreamThreadException(e, stream.name))
+            self._queue.put(PartitionGenerationCompletedSentinel(stream))

airbyte_cdk/sources/streams/concurrent/partition_reader.py

@@ -3,8 +3,12 @@
 #
 from queue import Queue
 
+from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.types import PartitionCompleteSentinel, QueueItem
+from airbyte_cdk.sources.streams.concurrent.partitions.types import (
+    PartitionCompleteSentinel,
+    QueueItem,
+)
 
 
 class PartitionReader:
@@ -12,6 +16,8 @@ class PartitionReader:
     Generates records from a partition and puts them in a queue.
     """
 
+    _IS_SUCCESSFUL = True
+
     def __init__(self, queue: Queue[QueueItem]) -> None:
         """
         :param queue: The queue to put the records in.
@@ -33,6 +39,7 @@ class PartitionReader:
         try:
             for record in partition.read():
                 self._queue.put(record)
-            self._queue.put(PartitionCompleteSentinel(partition))
+            self._queue.put(PartitionCompleteSentinel(partition, self._IS_SUCCESSFUL))
         except Exception as e:
-            self._queue.put(e)
+            self._queue.put(StreamThreadException(e, partition.stream_name()))
+            self._queue.put(PartitionCompleteSentinel(partition, not self._IS_SUCCESSFUL))

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

@@ -5,7 +5,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, Mapping, Optional
 
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.types import Record
 
 
 class Partition(ABC):
@@ -40,21 +40,6 @@ class Partition(ABC):
         """
         pass
 
-    @abstractmethod
-    def close(self) -> None:
-        """
-        Closes the partition.
-        """
-        pass
-
-    @abstractmethod
-    def is_closed(self) -> bool:
-        """
-        Returns whether the partition is closed.
-        :return:
-        """
-        pass
-
     @abstractmethod
     def __hash__(self) -> int:
         """