airbyte-cdk 0.0.0.dev0__py3-none-any.whl

airbyte_cdk/sources/streams/concurrent/clamping.py

@@ -0,0 +1,99 @@
+from abc import ABC
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import Callable
+
+from airbyte_cdk.sources.streams.concurrent.cursor_types import CursorValueType
+
+
+class ClampingStrategy(ABC):
+    def clamp(self, value: CursorValueType) -> CursorValueType:
+        raise NotImplementedError()
+
+
+class NoClamping(ClampingStrategy):
+    def clamp(self, value: CursorValueType) -> CursorValueType:
+        return value
+
+
+class ClampingEndProvider:
+    def __init__(
+        self,
+        clamping_strategy: ClampingStrategy,
+        end_provider: Callable[[], CursorValueType],
+        granularity: timedelta,
+    ) -> None:
+        self._clamping_strategy = clamping_strategy
+        self._end_provider = end_provider
+        self._granularity = granularity
+
+    def __call__(self) -> CursorValueType:
+        return self._clamping_strategy.clamp(self._end_provider()) - self._granularity
+
+
+class DayClampingStrategy(ClampingStrategy):
+    def __init__(self, is_ceiling: bool = True) -> None:
+        self._is_ceiling = is_ceiling
+
+    def clamp(self, value: datetime) -> datetime:  # type: ignore  # datetime implements method from CursorValueType
+        return_value = value.replace(hour=0, minute=0, second=0, microsecond=0)
+        if self._is_ceiling:
+            return return_value + timedelta(days=1)
+        return return_value
+
+
+class MonthClampingStrategy(ClampingStrategy):
+    def __init__(self, is_ceiling: bool = True) -> None:
+        self._is_ceiling = is_ceiling
+
+    def clamp(self, value: datetime) -> datetime:  # type: ignore  # datetime implements method from CursorValueType
+        return_value = value.replace(hour=0, minute=0, second=0, microsecond=0)
+        needs_to_round = value.day != 1
+        if not needs_to_round:
+            return return_value
+
+        return self._ceil(return_value) if self._is_ceiling else return_value.replace(day=1)
+
+    def _ceil(self, value: datetime) -> datetime:
+        return value.replace(
+            year=value.year + 1 if value.month == 12 else value.year,
+            month=(value.month % 12) + 1,
+            day=1,
+            hour=0,
+            minute=0,
+            second=0,
+            microsecond=0,
+        )
+
+
+class Weekday(Enum):
+    """
+    These integer values map to the same ones used by the Datetime.date.weekday() implementation
+    """
+
+    MONDAY = 0
+    TUESDAY = 1
+    WEDNESDAY = 2
+    THURSDAY = 3
+    FRIDAY = 4
+    SATURDAY = 5
+    SUNDAY = 6
+
+
+class WeekClampingStrategy(ClampingStrategy):
+    def __init__(self, day_of_week: Weekday, is_ceiling: bool = True) -> None:
+        self._day_of_week = day_of_week.value
+        self._is_ceiling = is_ceiling
+
+    def clamp(self, value: datetime) -> datetime:  # type: ignore  # datetime implements method from CursorValueType
+        days_diff_to_ceiling = (
+            7 - (value.weekday() - self._day_of_week)
+            if value.weekday() > self._day_of_week
+            else abs(value.weekday() - self._day_of_week)
+        )
+        delta = (
+            timedelta(days_diff_to_ceiling)
+            if self._is_ceiling
+            else timedelta(days_diff_to_ceiling - 7)
+        )
+        return value.replace(hour=0, minute=0, second=0, microsecond=0) + delta

airbyte_cdk/sources/streams/concurrent/cursor.py

@@ -0,0 +1,481 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import functools
+import logging
+from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    List,
+    Mapping,
+    MutableMapping,
+    Optional,
+    Tuple,
+    Union,
+)
+
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
+from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.streams import NO_CURSOR_STATE_KEY
+from airbyte_cdk.sources.streams.concurrent.clamping import ClampingStrategy, NoClamping
+from airbyte_cdk.sources.streams.concurrent.cursor_types import CursorValueType, GapType
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
+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 CursorField:
+    def __init__(self, cursor_field_key: str) -> None:
+        self.cursor_field_key = cursor_field_key
+
+    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(StreamSlicer, ABC):
+    @property
+    @abstractmethod
+    def state(self) -> MutableMapping[str, Any]: ...
+
+    @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()
+
+    @abstractmethod
+    def ensure_at_least_one_state_emitted(self) -> None:
+        """
+        State messages are emitted when a partition is closed. However, the platform expects at least one state to be emitted per sync per
+        stream. Hence, if no partitions are generated, this method needs to be called.
+        """
+        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."""
+
+    def __init__(
+        self,
+        stream_name: str,
+        stream_namespace: Optional[str],
+        message_repository: MessageRepository,
+    ) -> None:
+        self._stream_name = stream_name
+        self._stream_namespace = stream_namespace
+        self._message_repository = message_repository
+        # 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()
+        self._has_closed_at_least_one_slice = False
+
+    @property
+    def state(self) -> MutableMapping[str, Any]:
+        return {NO_CURSOR_STATE_KEY: True}
+
+    def observe(self, record: Record) -> None:
+        pass
+
+    def close_partition(self, partition: Partition) -> None:
+        pass
+
+    def ensure_at_least_one_state_emitted(self) -> None:
+        """
+        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._message_repository.emit_message(state_message)
+
+
+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,
+        connector_state_converter: AbstractStreamStateConverter,
+        cursor_field: CursorField,
+        slice_boundary_fields: Optional[Tuple[str, str]],
+        start: Optional[CursorValueType],
+        end_provider: Callable[[], CursorValueType],
+        lookback_window: Optional[GapType] = None,
+        slice_range: Optional[GapType] = None,
+        cursor_granularity: Optional[GapType] = None,
+        clamping_strategy: ClampingStrategy = NoClamping(),
+    ) -> None:
+        self._stream_name = stream_name
+        self._stream_namespace = stream_namespace
+        self._message_repository = message_repository
+        self._connector_state_converter = connector_state_converter
+        self._connector_state_manager = connector_state_manager
+        self._cursor_field = cursor_field
+        # To see some example where the slice boundaries might not be defined, check https://github.com/airbytehq/airbyte/blob/1ce84d6396e446e1ac2377362446e3fb94509461/airbyte-integrations/connectors/source-stripe/source_stripe/streams.py#L363-L379
+        self._slice_boundary_fields = slice_boundary_fields
+        self._start = start
+        self._end_provider = end_provider
+        self.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
+        self._clamping_strategy = clamping_strategy
+
+    @property
+    def state(self) -> MutableMapping[str, Any]:
+        return self._connector_state_converter.convert_to_state_message(
+            self.cursor_field, self._concurrent_state
+        )
+
+    @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
+        )
+
+    def observe(self, record: Record) -> None:
+        most_recent_cursor_value = self._most_recent_cursor_value_per_partition.get(
+            record.associated_slice
+        )
+        try:
+            cursor_value = self._extract_cursor_value(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))
+
+    def close_partition(self, partition: Partition) -> None:
+        slice_count_before = len(self._concurrent_state.get("slices", []))
+        self._add_slice_to_state(partition)
+        if slice_count_before < len(
+            self._concurrent_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._concurrent_state:
+                raise RuntimeError(
+                    f"The state for stream {self._stream_name} should have at least one slice to delineate the sync start time, but no slices are present. This is unexpected. Please contact Support."
+                )
+            self._concurrent_state["slices"].append(
+                {
+                    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 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
+                # boundaries. There are cases where partitions could not have boundaries:
+                # * The cursor should be per-partition
+                # * The stream state is actually the parent stream state
+                # There might be other cases not listed above. Those are not supported today hence the stream should not use this cursor for
+                # state management. For the specific user that was affected with this issue, we need to:
+                # * Fix state tracking (which is currently broken)
+                # * Make the new version available
+                # * (Probably) ask the user to reset the stream to avoid data loss
+                raise ValueError(
+                    "Given that slice_boundary_fields is not defined and that per-partition state is not supported, only one slice is "
+                    "expected. Please contact the Airbyte team."
+                )
+
+            self._concurrent_state["slices"].append(
+                {
+                    self._connector_state_converter.START_KEY: self.start,
+                    self._connector_state_converter.END_KEY: most_recent_cursor_value,
+                    self._connector_state_converter.MOST_RECENT_RECORD_KEY: most_recent_cursor_value,
+                }
+            )
+
+    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
+        )
+        self._message_repository.emit_message(state_message)
+
+    def _merge_partitions(self) -> None:
+        self._concurrent_state["slices"] = self._connector_state_converter.merge_intervals(
+            self._concurrent_state["slices"]
+        )
+
+    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 CursorValueType
+        except KeyError as 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:
+        """
+        The platform expect to have at least one state message on successful syncs. Hence, whatever happens, we expect this method to be
+        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._concurrent_state["slices"][0][self._connector_state_converter.START_KEY],
+                False,
+            )
+
+        if len(self._concurrent_state["slices"]) == 1:
+            yield from self._split_per_slice_range(
+                self._calculate_lower_boundary_of_last_slice(
+                    self._concurrent_state["slices"][0][self._connector_state_converter.END_KEY]
+                ),
+                self._end_provider(),
+                True,
+            )
+        elif len(self._concurrent_state["slices"]) > 1:
+            for i in range(len(self._concurrent_state["slices"]) - 1):
+                if self._cursor_granularity:
+                    yield from self._split_per_slice_range(
+                        self._concurrent_state["slices"][i][self._connector_state_converter.END_KEY]
+                        + self._cursor_granularity,
+                        self._concurrent_state["slices"][i + 1][
+                            self._connector_state_converter.START_KEY
+                        ],
+                        False,
+                    )
+                else:
+                    yield from self._split_per_slice_range(
+                        self._concurrent_state["slices"][i][
+                            self._connector_state_converter.END_KEY
+                        ],
+                        self._concurrent_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._concurrent_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._concurrent_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:
+            clamped_lower = self._clamping_strategy.clamp(lower)
+            clamped_upper = self._clamping_strategy.clamp(upper)
+            start_value, end_value = (
+                (clamped_lower, clamped_upper - self._cursor_granularity)
+                if self._cursor_granularity and not upper_is_end
+                else (clamped_lower, clamped_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
+
+                clamped_upper = (
+                    self._clamping_strategy.clamp(current_upper_boundary)
+                    if current_upper_boundary != upper
+                    else current_upper_boundary
+                )
+                clamped_lower = self._clamping_strategy.clamp(current_lower_boundary)
+                if clamped_lower >= clamped_upper:
+                    # clamping collapsed both values which means that it is time to stop processing
+                    # FIXME should this be replace by proper end_provider
+                    break
+                start_value, end_value = (
+                    (clamped_lower, clamped_upper - self._cursor_granularity)
+                    if self._cursor_granularity
+                    and (not upper_is_end or not has_reached_upper_boundary)
+                    else (clamped_lower, clamped_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),
+                    },
+                )
+                current_lower_boundary = clamped_upper
+                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/cursor_types.py

@@ -0,0 +1,32 @@
+from abc import abstractmethod
+from typing import 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: "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