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

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

@@ -0,0 +1,21 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from abc import ABC, abstractmethod
+from typing import Iterable
+
+from airbyte_cdk.sources.types import StreamSlice
+
+
+class StreamSlicer(ABC):
+    """
+    Slices the stream into chunks that can be fetched independently. Slices enable state checkpointing and data retrieval parallelization.
+    """
+
+    @abstractmethod
+    def stream_slices(self) -> Iterable[StreamSlice]:
+        """
+        Defines stream slices
+
+        :return: An iterable of stream slices
+        """
+        pass

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

@@ -2,11 +2,13 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from typing import Union
+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
+from airbyte_cdk.sources.types import Record
 
 
 class PartitionCompleteSentinel:
@@ -15,14 +17,22 @@ class PartitionCompleteSentinel:
     Includes a pointer to the partition that was processed.
     """
 
-    def __init__(self, partition: Partition):
+    def __init__(self, partition: Partition, is_successful: bool = True):
         """
         :param partition: The partition that was processed
         """
         self.partition = partition
+        self.is_successful = is_successful
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, PartitionCompleteSentinel):
+            return self.partition == other.partition
+        return False
 
 
 """
 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

@@ -4,7 +4,7 @@
 
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import TYPE_CHECKING, Any, List, MutableMapping, Tuple
+from typing import TYPE_CHECKING, Any, List, MutableMapping, Optional, Tuple
 
 if TYPE_CHECKING:
     from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
@@ -17,24 +17,92 @@ class ConcurrencyCompatibleStateType(Enum):
 class AbstractStreamStateConverter(ABC):
     START_KEY = "start"
     END_KEY = "end"
+    MOST_RECENT_RECORD_KEY = "most_recent_cursor_value"
 
     @abstractmethod
+    def _from_state_message(self, value: Any) -> Any:
+        pass
+
+    @abstractmethod
+    def _to_state_message(self, value: Any) -> Any:
+        pass
+
+    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]:
+        """
+        Convert the state message from the concurrency-compatible format to the stream's original format.
+
+        e.g.
+        { "created": "2021-01-18T21:18:20.000Z" }
+        """
+        if self.is_state_message_compatible(stream_state) and self._is_sequential_state:
+            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)}
+                )
+            return legacy_state or {}
+        else:
+            return self.serialize(stream_state, ConcurrencyCompatibleStateType.date_range)
+
+    def _get_latest_complete_time(self, slices: List[MutableMapping[str, Any]]) -> Any:
+        """
+        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."
+            )
+        merged_intervals = self.merge_intervals(slices)
+        first_interval = merged_intervals[0]
+
+        return first_interval.get("most_recent_cursor_value") or first_interval[self.START_KEY]
+
     def deserialize(self, state: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
         """
         Perform any transformations needed for compatibility with the converter.
         """
-        ...
+        for stream_slice in state.get("slices", []):
+            stream_slice[self.START_KEY] = self._from_state_message(stream_slice[self.START_KEY])
+            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]:
+        """
+        Perform any transformations needed for compatibility with the converter.
+        """
+        serialized_slices = []
+        for stream_slice in state.get("slices", []):
+            serialized_slice = {
+                self.START_KEY: self._to_state_message(stream_slice[self.START_KEY]),
+                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_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(
         self,
-        cursor_field: "CursorField",
+        cursor_field: "CursorField",  # to deprecate as it is only needed for sequential state
         stream_state: MutableMapping[str, Any],
-        start: Any,
+        start: Optional[Any],
     ) -> Tuple[Any, MutableMapping[str, Any]]:
         """
         Convert the state message to the format required by the ConcurrentCursor.
@@ -50,31 +118,56 @@ class AbstractStreamStateConverter(ABC):
         ...
 
     @abstractmethod
-    def convert_to_sequential_state(self, cursor_field: "CursorField", stream_state: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
+    def increment(self, value: Any) -> Any:
         """
-        Convert the state message from the concurrency-compatible format to the stream's original format.
-
-        e.g.
-        { "created": 1617030403 }
+        Increment a timestamp by a single unit.
         """
         ...
 
     @abstractmethod
-    def increment(self, timestamp: Any) -> Any:
+    def output_format(self, value: Any) -> Any:
         """
-        Increment a timestamp by a single unit.
+        Convert the cursor value type to a JSON valid type.
         """
         ...
 
-    @abstractmethod
-    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.
 
         Intervals may be merged if the start time of the second interval is 1 unit or less (as defined by the
         `increment` method) than the end time of the first interval.
         """
-        ...
+        if not intervals:
+            return []
+
+        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:]:
+            last_interval = merged_intervals[-1]
+            last_interval_end = last_interval[self.END_KEY]
+            current_interval_start = current_interval[self.START_KEY]
+
+            if self.increment(last_interval_end) >= current_interval_start:
+                last_interval[self.END_KEY] = max(last_interval_end, current_interval[self.END_KEY])
+                last_interval_cursor_value = last_interval.get("most_recent_cursor_value")
+                current_interval_cursor_value = current_interval.get("most_recent_cursor_value")
+
+                last_interval["most_recent_cursor_value"] = (
+                    max(current_interval_cursor_value, last_interval_cursor_value)
+                    if current_interval_cursor_value and last_interval_cursor_value
+                    else current_interval_cursor_value or last_interval_cursor_value
+                )
+            else:
+                # Add a new interval if no overlap
+                merged_intervals.append(current_interval)
+
+        return merged_intervals
 
     @abstractmethod
     def parse_value(self, value: Any) -> Any:
@@ -85,5 +178,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

@@ -3,45 +3,49 @@
 #
 
 from abc import abstractmethod
-from datetime import datetime, timedelta
-from typing import Any, List, MutableMapping, Optional, Tuple
+from datetime import datetime, timedelta, timezone
+from typing import Any, Callable, List, MutableMapping, Optional, Tuple
 
 import pendulum
+from pendulum.datetime import DateTime
+
+# FIXME We would eventually like the Concurrent package do be agnostic of the declarative package. However, this is a breaking change and
+#  the goal in the short term is only to fix the issue we are seeing for source-declarative-manifest.
+from airbyte_cdk.sources.declarative.datetime.datetime_parser import DatetimeParser
 from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
 from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter import (
     AbstractStreamStateConverter,
     ConcurrencyCompatibleStateType,
 )
-from pendulum.datetime import DateTime
 
 
 class DateTimeStreamStateConverter(AbstractStreamStateConverter):
+    def _from_state_message(self, value: Any) -> Any:
+        return self.parse_timestamp(value)
+
+    def _to_state_message(self, value: Any) -> Any:
+        return self.output_format(value)
+
     @property
     @abstractmethod
-    def _zero_value(self) -> Any:
-        ...
+    def _zero_value(self) -> Any: ...
 
     @property
     def zero_value(self) -> datetime:
         return self.parse_timestamp(self._zero_value)
 
-    @abstractmethod
-    def increment(self, timestamp: datetime) -> datetime:
-        ...
+    @classmethod
+    def get_end_provider(cls) -> Callable[[], datetime]:
+        return lambda: datetime.now(timezone.utc)
 
     @abstractmethod
-    def parse_timestamp(self, timestamp: Any) -> datetime:
-        ...
+    def increment(self, timestamp: datetime) -> datetime: ...
 
     @abstractmethod
-    def output_format(self, timestamp: datetime) -> Any:
-        ...
+    def parse_timestamp(self, timestamp: Any) -> datetime: ...
 
-    def deserialize(self, state: MutableMapping[str, Any]) -> MutableMapping[str, Any]:
-        for stream_slice in state.get("slices", []):
-            stream_slice[self.START_KEY] = self.parse_timestamp(stream_slice[self.START_KEY])
-            stream_slice[self.END_KEY] = self.parse_timestamp(stream_slice[self.END_KEY])
-        return state
+    @abstractmethod
+    def output_format(self, timestamp: datetime) -> Any: ...
 
     def parse_value(self, value: Any) -> Any:
         """
@@ -49,29 +53,14 @@ class DateTimeStreamStateConverter(AbstractStreamStateConverter):
         """
         return self.parse_timestamp(value)
 
-    def merge_intervals(self, intervals: List[MutableMapping[str, datetime]]) -> List[MutableMapping[str, datetime]]:
-        if not intervals:
-            return []
-
-        sorted_intervals = sorted(intervals, key=lambda x: (x[self.START_KEY], x[self.END_KEY]))
-        merged_intervals = [sorted_intervals[0]]
-
-        for interval in sorted_intervals[1:]:
-            last_end_time = merged_intervals[-1][self.END_KEY]
-            current_start_time = interval[self.START_KEY]
-            if self._compare_intervals(last_end_time, current_start_time):
-                merged_end_time = max(last_end_time, interval[self.END_KEY])
-                merged_intervals[-1][self.END_KEY] = merged_end_time
-            else:
-                merged_intervals.append(interval)
-
-        return merged_intervals
-
     def _compare_intervals(self, end_time: Any, start_time: Any) -> bool:
         return bool(self.increment(end_time) >= start_time)
 
     def convert_from_sequential_state(
-        self, cursor_field: CursorField, stream_state: MutableMapping[str, Any], start: 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.
@@ -92,7 +81,13 @@ 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: 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,
+                self.MOST_RECENT_RECORD_KEY: sync_start,
+            }
+        ]
 
         return sync_start, {
             "state_type": ConcurrencyCompatibleStateType.date_range.value,
@@ -100,43 +95,23 @@ class DateTimeStreamStateConverter(AbstractStreamStateConverter):
             "legacy": stream_state,
         }
 
-    def _get_sync_start(self, cursor_field: CursorField, stream_state: MutableMapping[str, Any], start: Optional[Any]) -> datetime:
-        sync_start = self.parse_timestamp(start) if start is not None else self.zero_value
+    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
         else:
             return sync_start
 
-    def convert_to_sequential_state(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.
-
-        e.g.
-        { "created": "2021-01-18T21:18:20.000Z" }
-        """
-        if self.is_state_message_compatible(stream_state):
-            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.output_format(latest_complete_time)})
-            return legacy_state or {}
-        else:
-            return stream_state
-
-    def _get_latest_complete_time(self, slices: List[MutableMapping[str, Any]]) -> Optional[datetime]:
-        """
-        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.")
-
-        merged_intervals = self.merge_intervals(slices)
-        first_interval = merged_intervals[0]
-        return first_interval[self.END_KEY]
-
 
 class EpochValueConcurrentStreamStateConverter(DateTimeStreamStateConverter):
     """
@@ -163,8 +138,10 @@ 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})")
-        return dt_object  # type: ignore  # we are manually type checking because pendulum.parse may return different types
+            raise ValueError(
+                f"DateTime object was expected but got {type(dt_object)} from pendulum.parse({timestamp})"
+            )
+        return dt_object
 
 
 class IsoMillisConcurrentStreamStateConverter(DateTimeStreamStateConverter):
@@ -183,8 +160,14 @@ class IsoMillisConcurrentStreamStateConverter(DateTimeStreamStateConverter):
 
     _zero_value = "0001-01-01T00:00:00.000Z"
 
+    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)
+
     def increment(self, timestamp: datetime) -> datetime:
-        return timestamp + timedelta(milliseconds=1)
+        return timestamp + self._cursor_granularity
 
     def output_format(self, timestamp: datetime) -> Any:
         return timestamp.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
@@ -192,5 +175,40 @@ 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})")
-        return dt_object  # type: ignore  # we are manually type checking because pendulum.parse may return different types
+            raise ValueError(
+                f"DateTime object was expected but got {type(dt_object)} from pendulum.parse({timestamp})"
+            )
+        return dt_object
+
+
+class CustomFormatConcurrentStreamStateConverter(IsoMillisConcurrentStreamStateConverter):
+    """
+    Datetime State converter that emits state according to the supplied datetime format. The converter supports reading
+    incoming state in any valid datetime format via Pendulum.
+    """
+
+    def __init__(
+        self,
+        datetime_format: str,
+        input_datetime_formats: Optional[List[str]] = None,
+        is_sequential_state: bool = True,
+        cursor_granularity: Optional[timedelta] = None,
+    ):
+        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]
+        self._parser = DatetimeParser()
+
+    def output_format(self, timestamp: datetime) -> str:
+        return self._parser.format(timestamp, self._datetime_format)
+
+    def parse_timestamp(self, timestamp: str) -> datetime:
+        for datetime_format in self._input_datetime_formats:
+            try:
+                return self._parser.parse(timestamp, datetime_format)
+            except ValueError:
+                pass
+        raise ValueError(f"No format in {self._input_datetime_formats} matching {timestamp}")