airbyte-cdk 6.5.3rc2__py3-none-any.whl → 6.6.0__py3-none-any.whl

airbyte_cdk/sources/streams/concurrent/cursor.py

@@ -11,7 +11,9 @@ 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.state_converters.abstract_stream_state_converter import AbstractStreamStateConverter
+from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter import (
+    AbstractStreamStateConverter,
+)
 
 
 def _extract_value(mapping: Mapping[str, Any], path: List[str]) -> Any:
@@ -62,8 +64,7 @@ class CursorField:
 class Cursor(ABC):
     @property
     @abstractmethod
-    def state(self) -> MutableMapping[str, Any]:
-        ...
+    def state(self) -> MutableMapping[str, Any]: ...
 
     @abstractmethod
     def observe(self, record: Record) -> None:
@@ -128,8 +129,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)
 
 
@@ -182,13 +187,22 @@ class ConcurrentCursor(Cursor):
     def slice_boundary_fields(self) -> Optional[Tuple[str, str]]:
         return self._slice_boundary_fields
 
-    def _get_concurrent_state(self, state: MutableMapping[str, Any]) -> Tuple[CursorValueType, MutableMapping[str, Any]]:
+    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:
-        most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(record.partition)
+        most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(
+            record.partition
+        )
         cursor_value = self._extract_cursor_value(record)
 
         if most_recent_cursor_value is None or most_recent_cursor_value < cursor_value:
@@ -200,7 +214,9 @@ 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
@@ -253,9 +269,13 @@ class ConcurrentCursor(Cursor):
         self._connector_state_manager.update_state_for_stream(
             self._stream_name,
             self._stream_namespace,
-            self._connector_state_converter.convert_to_state_message(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
         )
-        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:
@@ -268,7 +288,9 @@ class ConcurrentCursor(Cursor):
                 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 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:
         """
@@ -300,7 +322,9 @@ class ConcurrentCursor(Cursor):
 
         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._calculate_lower_boundary_of_last_slice(
+                    self.state["slices"][0][self._connector_state_converter.END_KEY]
+                ),
                 self._end_provider(),
                 True,
             )
@@ -308,7 +332,8 @@ class ConcurrentCursor(Cursor):
             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][self._connector_state_converter.END_KEY]
+                        + self._cursor_granularity,
                         self.state["slices"][i + 1][self._connector_state_converter.START_KEY],
                         False,
                     )
@@ -319,7 +344,9 @@ class ConcurrentCursor(Cursor):
                         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._calculate_lower_boundary_of_last_slice(
+                    self.state["slices"][-1][self._connector_state_converter.END_KEY]
+                ),
                 self._end_provider(),
                 True,
             )
@@ -327,9 +354,14 @@ class ConcurrentCursor(Cursor):
             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]
+        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:
+    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
@@ -353,9 +385,13 @@ class ConcurrentCursor(Cursor):
             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)
+                current_upper_boundary = min(
+                    self._evaluate_upper_safely(current_lower_boundary, self._slice_range), upper
+                )
                 has_reached_upper_boundary = current_upper_boundary >= upper
-                if self._cursor_granularity and (not upper_is_end or not has_reached_upper_boundary):
+                if self._cursor_granularity and (
+                    not upper_is_end or not has_reached_upper_boundary
+                ):
                     yield current_lower_boundary, current_upper_boundary - self._cursor_granularity
                 else:
                     yield current_lower_boundary, current_upper_boundary

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,7 +63,11 @@ 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],
+        )
 
         if self._namespace:
             stream.namespace = self._namespace

airbyte_cdk/sources/streams/concurrent/helpers.py

@@ -5,7 +5,9 @@ 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):
@@ -22,7 +24,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,9 @@
 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
@@ -16,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.

airbyte_cdk/sources/streams/concurrent/partition_reader.py

@@ -5,7 +5,10 @@ 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:

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

@@ -13,7 +13,12 @@ class Record:
     Represents a record read from a stream.
     """
 
-    def __init__(self, data: Mapping[str, Any], partition: "Partition", is_file_transfer_message: bool = False):
+    def __init__(
+        self,
+        data: Mapping[str, Any],
+        partition: "Partition",
+        is_file_transfer_message: bool = False,
+    ):
         self.data = data
         self.partition = partition
         self.is_file_transfer_message = is_file_transfer_message
@@ -21,7 +26,10 @@ class Record:
     def __eq__(self, other: Any) -> bool:
         if not isinstance(other, Record):
             return False
-        return self.data == other.data and self.partition.stream_name() == other.partition.stream_name()
+        return (
+            self.data == other.data
+            and self.partition.stream_name() == other.partition.stream_name()
+        )
 
     def __repr__(self) -> str:
         return f"Record(data={self.data}, stream_name={self.partition.stream_name()})"

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

@@ -4,7 +4,9 @@
 
 from typing import Any, Union
 
-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.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 
@@ -31,4 +33,6 @@ class PartitionCompleteSentinel:
 """
 Typedef representing the items that can be added to the ThreadBasedConcurrentStream
 """
-QueueItem = Union[Record, Partition, PartitionCompleteSentinel, PartitionGenerationCompletedSentinel, Exception]
+QueueItem = Union[
+    Record, Partition, PartitionCompleteSentinel, PartitionGenerationCompletedSentinel, Exception
+]

airbyte_cdk/sources/streams/concurrent/state_converters/abstract_stream_state_converter.py

@@ -30,7 +30,9 @@ class AbstractStreamStateConverter(ABC):
     def __init__(self, is_sequential_state: bool = True):
         self._is_sequential_state = is_sequential_state
 
-    def convert_to_state_message(self, cursor_field: "CursorField", stream_state: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
+    def convert_to_state_message(
+        self, cursor_field: "CursorField", stream_state: MutableMapping[str, Any]
+    ) -> MutableMapping[str, Any]:
         """
         Convert the state message from the concurrency-compatible format to the stream's original format.
 
@@ -41,7 +43,9 @@ class AbstractStreamStateConverter(ABC):
             legacy_state = stream_state.get("legacy", {})
             latest_complete_time = self._get_latest_complete_time(stream_state.get("slices", []))
             if latest_complete_time is not None:
-                legacy_state.update({cursor_field.cursor_field_key: self._to_state_message(latest_complete_time)})
+                legacy_state.update(
+                    {cursor_field.cursor_field_key: self._to_state_message(latest_complete_time)}
+                )
             return legacy_state or {}
         else:
             return self.serialize(stream_state, ConcurrencyCompatibleStateType.date_range)
@@ -51,7 +55,9 @@ class AbstractStreamStateConverter(ABC):
         Get the latest time before which all records have been processed.
         """
         if not slices:
-            raise RuntimeError("Expected at least one slice but there were none. This is unexpected; please contact Support.")
+            raise RuntimeError(
+                "Expected at least one slice but there were none. This is unexpected; please contact Support."
+            )
         merged_intervals = self.merge_intervals(slices)
         first_interval = merged_intervals[0]
 
@@ -66,7 +72,9 @@ class AbstractStreamStateConverter(ABC):
             stream_slice[self.END_KEY] = self._from_state_message(stream_slice[self.END_KEY])
         return state
 
-    def serialize(self, state: MutableMapping[str, Any], state_type: ConcurrencyCompatibleStateType) -> MutableMapping[str, Any]:
+    def serialize(
+        self, state: MutableMapping[str, Any], state_type: ConcurrencyCompatibleStateType
+    ) -> MutableMapping[str, Any]:
         """
         Perform any transformations needed for compatibility with the converter.
         """
@@ -77,13 +85,17 @@ class AbstractStreamStateConverter(ABC):
                 self.END_KEY: self._to_state_message(stream_slice[self.END_KEY]),
             }
             if stream_slice.get(self.MOST_RECENT_RECORD_KEY):
-                serialized_slice[self.MOST_RECENT_RECORD_KEY] = self._to_state_message(stream_slice[self.MOST_RECENT_RECORD_KEY])
+                serialized_slice[self.MOST_RECENT_RECORD_KEY] = self._to_state_message(
+                    stream_slice[self.MOST_RECENT_RECORD_KEY]
+                )
             serialized_slices.append(serialized_slice)
         return {"slices": serialized_slices, "state_type": state_type.value}
 
     @staticmethod
     def is_state_message_compatible(state: MutableMapping[str, Any]) -> bool:
-        return bool(state) and state.get("state_type") in [t.value for t in ConcurrencyCompatibleStateType]
+        return bool(state) and state.get("state_type") in [
+            t.value for t in ConcurrencyCompatibleStateType
+        ]
 
     @abstractmethod
     def convert_from_sequential_state(
@@ -112,7 +124,9 @@ class AbstractStreamStateConverter(ABC):
         """
         ...
 
-    def merge_intervals(self, intervals: List[MutableMapping[str, Any]]) -> List[MutableMapping[str, Any]]:
+    def merge_intervals(
+        self, intervals: List[MutableMapping[str, Any]]
+    ) -> List[MutableMapping[str, Any]]:
         """
         Compute and return a list of merged intervals.
 
@@ -122,7 +136,9 @@ class AbstractStreamStateConverter(ABC):
         if not intervals:
             return []
 
-        sorted_intervals = sorted(intervals, key=lambda interval: (interval[self.START_KEY], interval[self.END_KEY]))
+        sorted_intervals = sorted(
+            intervals, key=lambda interval: (interval[self.START_KEY], interval[self.END_KEY])
+        )
         merged_intervals = [sorted_intervals[0]]
 
         for current_interval in sorted_intervals[1:]:
@@ -155,5 +171,4 @@ class AbstractStreamStateConverter(ABC):
 
     @property
     @abstractmethod
-    def zero_value(self) -> Any:
-        ...
+    def zero_value(self) -> Any: ...

airbyte_cdk/sources/streams/concurrent/state_converters/datetime_stream_state_converter.py

@@ -28,8 +28,7 @@ class DateTimeStreamStateConverter(AbstractStreamStateConverter):
 
     @property
     @abstractmethod
-    def _zero_value(self) -> Any:
-        ...
+    def _zero_value(self) -> Any: ...
 
     @property
     def zero_value(self) -> datetime:
@@ -40,16 +39,13 @@ class DateTimeStreamStateConverter(AbstractStreamStateConverter):
         return lambda: datetime.now(timezone.utc)
 
     @abstractmethod
-    def increment(self, timestamp: datetime) -> datetime:
-        ...
+    def increment(self, timestamp: datetime) -> datetime: ...
 
     @abstractmethod
-    def parse_timestamp(self, timestamp: Any) -> datetime:
-        ...
+    def parse_timestamp(self, timestamp: Any) -> datetime: ...
 
     @abstractmethod
-    def output_format(self, timestamp: datetime) -> Any:
-        ...
+    def output_format(self, timestamp: datetime) -> Any: ...
 
     def parse_value(self, value: Any) -> Any:
         """
@@ -61,7 +57,10 @@ class DateTimeStreamStateConverter(AbstractStreamStateConverter):
         return bool(self.increment(end_time) >= start_time)
 
     def convert_from_sequential_state(
-        self, cursor_field: CursorField, stream_state: MutableMapping[str, Any], start: Optional[datetime]
+        self,
+        cursor_field: CursorField,
+        stream_state: MutableMapping[str, Any],
+        start: Optional[datetime],
     ) -> Tuple[datetime, MutableMapping[str, Any]]:
         """
         Convert the state message to the format required by the ConcurrentCursor.
@@ -82,7 +81,9 @@ class DateTimeStreamStateConverter(AbstractStreamStateConverter):
         # Create a slice to represent the records synced during prior syncs.
         # The start and end are the same to avoid confusion as to whether the records for this slice
         # were actually synced
-        slices = [{self.START_KEY: start if start is not None else sync_start, self.END_KEY: sync_start}]
+        slices = [
+            {self.START_KEY: start if start is not None else sync_start, self.END_KEY: sync_start}
+        ]
 
         return sync_start, {
             "state_type": ConcurrencyCompatibleStateType.date_range.value,
@@ -90,10 +91,17 @@ class DateTimeStreamStateConverter(AbstractStreamStateConverter):
             "legacy": stream_state,
         }
 
-    def _get_sync_start(self, cursor_field: CursorField, stream_state: MutableMapping[str, Any], start: Optional[datetime]) -> datetime:
+    def _get_sync_start(
+        self,
+        cursor_field: CursorField,
+        stream_state: MutableMapping[str, Any],
+        start: Optional[datetime],
+    ) -> datetime:
         sync_start = start if start is not None else self.zero_value
         prev_sync_low_water_mark = (
-            self.parse_timestamp(stream_state[cursor_field.cursor_field_key]) if cursor_field.cursor_field_key in stream_state else None
+            self.parse_timestamp(stream_state[cursor_field.cursor_field_key])
+            if cursor_field.cursor_field_key in stream_state
+            else None
         )
         if prev_sync_low_water_mark and prev_sync_low_water_mark >= sync_start:
             return prev_sync_low_water_mark
@@ -126,7 +134,9 @@ class EpochValueConcurrentStreamStateConverter(DateTimeStreamStateConverter):
     def parse_timestamp(self, timestamp: int) -> datetime:
         dt_object = pendulum.from_timestamp(timestamp)
         if not isinstance(dt_object, DateTime):
-            raise ValueError(f"DateTime object was expected but got {type(dt_object)} from pendulum.parse({timestamp})")
+            raise ValueError(
+                f"DateTime object was expected but got {type(dt_object)} from pendulum.parse({timestamp})"
+            )
         return dt_object  # type: ignore  # we are manually type checking because pendulum.parse may return different types
 
 
@@ -146,7 +156,9 @@ class IsoMillisConcurrentStreamStateConverter(DateTimeStreamStateConverter):
 
     _zero_value = "0001-01-01T00:00:00.000Z"
 
-    def __init__(self, is_sequential_state: bool = True, cursor_granularity: Optional[timedelta] = None):
+    def __init__(
+        self, is_sequential_state: bool = True, cursor_granularity: Optional[timedelta] = None
+    ):
         super().__init__(is_sequential_state=is_sequential_state)
         self._cursor_granularity = cursor_granularity or timedelta(milliseconds=1)
 
@@ -159,7 +171,9 @@ class IsoMillisConcurrentStreamStateConverter(DateTimeStreamStateConverter):
     def parse_timestamp(self, timestamp: str) -> datetime:
         dt_object = pendulum.parse(timestamp)
         if not isinstance(dt_object, DateTime):
-            raise ValueError(f"DateTime object was expected but got {type(dt_object)} from pendulum.parse({timestamp})")
+            raise ValueError(
+                f"DateTime object was expected but got {type(dt_object)} from pendulum.parse({timestamp})"
+            )
         return dt_object  # type: ignore  # we are manually type checking because pendulum.parse may return different types
 
 
@@ -176,7 +190,9 @@ class CustomFormatConcurrentStreamStateConverter(IsoMillisConcurrentStreamStateC
         is_sequential_state: bool = True,
         cursor_granularity: Optional[timedelta] = None,
     ):
-        super().__init__(is_sequential_state=is_sequential_state, cursor_granularity=cursor_granularity)
+        super().__init__(
+            is_sequential_state=is_sequential_state, cursor_granularity=cursor_granularity
+        )
         self._datetime_format = datetime_format
         self._input_datetime_formats = input_datetime_formats if input_datetime_formats else []
         self._input_datetime_formats += [self._datetime_format]