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

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -0,0 +1,102 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from functools import lru_cache
+from logging import Logger
+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.cursor import Cursor
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
+from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
+
+
+class DefaultStream(AbstractStream):
+    def __init__(
+        self,
+        partition_generator: PartitionGenerator,
+        name: str,
+        json_schema: Mapping[str, Any],
+        availability_strategy: AbstractAvailabilityStrategy,
+        primary_key: List[str],
+        cursor_field: Optional[str],
+        logger: Logger,
+        cursor: Cursor,
+        namespace: Optional[str] = None,
+        supports_file_transfer: bool = False,
+    ) -> None:
+        self._stream_partition_generator = partition_generator
+        self._name = name
+        self._json_schema = json_schema
+        self._availability_strategy = availability_strategy
+        self._primary_key = primary_key
+        self._cursor_field = cursor_field
+        self._logger = logger
+        self._cursor = cursor
+        self._namespace = namespace
+        self._supports_file_transfer = supports_file_transfer
+
+    def generate_partitions(self) -> Iterable[Partition]:
+        yield from self._stream_partition_generator.generate()
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def namespace(self) -> Optional[str]:
+        return self._namespace
+
+    def check_availability(self) -> StreamAvailability:
+        return self._availability_strategy.check_availability(self._logger)
+
+    @property
+    def cursor_field(self) -> Optional[str]:
+        return self._cursor_field
+
+    @lru_cache(maxsize=None)
+    def get_json_schema(self) -> Mapping[str, Any]:
+        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],
+            is_resumable=False,
+            is_file_based=self._supports_file_transfer,
+        )
+
+        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 = [[key] for key in keys]
+
+        return stream
+
+    def log_stream_sync_configuration(self) -> None:
+        self._logger.debug(
+            f"Syncing stream instance: {self.name}",
+            extra={
+                "primary_key": self._primary_key,
+                "cursor_field": self.cursor_field,
+            },
+        )
+
+    @property
+    def cursor(self) -> Cursor:
+        return self._cursor

airbyte_cdk/sources/streams/concurrent/exceptions.py

@@ -0,0 +1,18 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from typing import Any
+
+
+class ExceptionWithDisplayMessage(Exception):
+    """
+    Exception that can be used to display a custom message to the user.
+    """
+
+    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

@@ -0,0 +1,42 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+
+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]:
+    if stream_primary_key is None:
+        return []
+    elif isinstance(stream_primary_key, str):
+        return [stream_primary_key]
+    elif isinstance(stream_primary_key, list):
+        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:
+        raise ValueError(f"Invalid type for primary key: {stream_primary_key}")
+
+
+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}"
+            )
+        elif len(stream.cursor_field) == 0:
+            return None
+        else:
+            return stream.cursor_field[0]
+    else:
+        return stream.cursor_field

airbyte_cdk/sources/streams/concurrent/partition_enqueuer.py

@@ -0,0 +1,64 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+import time
+from queue import Queue
+
+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
+
+
+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:
+        """
+        :param queue:  The queue to put the partitions in.
+        :param throttler: The throttler to use to throttle the partition generation.
+        """
+        self._queue = queue
+        self._thread_pool_manager = thread_pool_manager
+        self._sleep_time_in_seconds = sleep_time_in_seconds
+
+    def generate_partitions(self, stream: AbstractStream) -> None:
+        """
+        Generate partitions from a partition generator and put them in a queue.
+        When all the partitions are added to the queue, a sentinel is added to the queue to indicate that all the partitions have been generated.
+
+        If an exception is encountered, the exception will be caught and put in the queue. This is very important because if we don't, the
+        main thread will have no way to know that something when wrong and will wait until the timeout is reached
+
+        This method is meant to be called in a separate thread.
+        """
+        try:
+            for partition in stream.generate_partitions():
+                # Adding partitions to the queue generates futures. To avoid having too many futures, we throttle here. We understand that
+                # we might add more futures than the limit by throttling in the threads while it is the main thread that actual adds the
+                # future but we expect the delta between the max futures length and the actual to be small enough that it would not be an
+                # issue. We do this in the threads because we want the main thread to always be processing QueueItems as if it does not, the
+                # queue size could grow and generating OOM issues.
+                #
+                # Also note that we do not expect this to create deadlocks where all worker threads wait because we have less
+                # PartitionEnqueuer threads than worker threads.
+                #
+                # Also note that prune_to_validate_has_reached_futures_limit has a lock while pruning which might create a bottleneck in
+                # terms of performance.
+                while self._thread_pool_manager.prune_to_validate_has_reached_futures_limit():
+                    time.sleep(self._sleep_time_in_seconds)
+                self._queue.put(partition)
+            self._queue.put(PartitionGenerationCompletedSentinel(stream))
+        except Exception as e:
+            self._queue.put(StreamThreadException(e, stream.name))
+            self._queue.put(PartitionGenerationCompletedSentinel(stream))

airbyte_cdk/sources/streams/concurrent/partition_reader.py

@@ -0,0 +1,45 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+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,
+)
+
+
+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.
+        """
+        self._queue = queue
+
+    def process_partition(self, partition: Partition) -> None:
+        """
+        Process a partition and put the records in the output queue.
+        When all the partitions are added to the queue, a sentinel is added to the queue to indicate that all the partitions have been generated.
+
+        If an exception is encountered, the exception will be caught and put in the queue. This is very important because if we don't, the
+        main thread will have no way to know that something when wrong and will wait until the timeout is reached
+
+        This method is meant to be called from a thread.
+        :param partition: The partition to read data from
+        :return: None
+        """
+        try:
+            for record in partition.read():
+                self._queue.put(record)
+            self._queue.put(PartitionCompleteSentinel(partition, self._IS_SUCCESSFUL))
+        except Exception as e:
+            self._queue.put(StreamThreadException(e, partition.stream_name()))
+            self._queue.put(PartitionCompleteSentinel(partition, not self._IS_SUCCESSFUL))

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

@@ -0,0 +1,3 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#

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

@@ -0,0 +1,48 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from abc import ABC, abstractmethod
+from typing import Any, Iterable, Mapping, Optional
+
+from airbyte_cdk.sources.types import Record
+
+
+class Partition(ABC):
+    """
+    A partition is responsible for reading a specific set of data from a source.
+    """
+
+    @abstractmethod
+    def read(self) -> Iterable[Record]:
+        """
+        Reads the data from the partition.
+        :return: An iterable of records.
+        """
+        pass
+
+    @abstractmethod
+    def to_slice(self) -> Optional[Mapping[str, Any]]:
+        """
+        Converts the partition to a slice that can be serialized and deserialized.
+
+        Note: it would have been interesting to have a type of `Mapping[str, Comparable]` to simplify typing but some slices can have nested
+         values ([example](https://github.com/airbytehq/airbyte/blob/1ce84d6396e446e1ac2377362446e3fb94509461/airbyte-integrations/connectors/source-stripe/source_stripe/streams.py#L584-L596))
+        :return: A mapping representing a slice
+        """
+        pass
+
+    @abstractmethod
+    def stream_name(self) -> str:
+        """
+        Returns the name of the stream that this partition is reading from.
+        :return: The name of the stream.
+        """
+        pass
+
+    @abstractmethod
+    def __hash__(self) -> int:
+        """
+        Returns a hash of the partition.
+        Partitions must be hashable so that they can be used as keys in a dictionary.
+        """

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

@@ -0,0 +1,18 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from abc import ABC, abstractmethod
+from typing import Iterable
+
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
+
+
+class PartitionGenerator(ABC):
+    @abstractmethod
+    def generate(self) -> Iterable[Partition]:
+        """
+        Generates partitions for a given sync mode.
+        :return: An iterable of partitions
+        """
+        pass

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

@@ -0,0 +1,38 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from typing import Any, Union
+
+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.types import Record
+
+
+class PartitionCompleteSentinel:
+    """
+    A sentinel object indicating all records for a partition were produced.
+    Includes a pointer to the partition that was processed.
+    """
+
+    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
+]

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

@@ -0,0 +1,182 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Callable, List, MutableMapping, Optional, Tuple
+
+if TYPE_CHECKING:
+    from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
+
+
+class ConcurrencyCompatibleStateType(Enum):
+    date_range = "date-range"
+    integer = "integer"
+
+
+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
+        ]
+
+    @abstractmethod
+    def convert_from_sequential_state(
+        self,
+        cursor_field: "CursorField",  # to deprecate as it is only needed for sequential state
+        stream_state: MutableMapping[str, Any],
+        start: Optional[Any],
+    ) -> Tuple[Any, MutableMapping[str, Any]]:
+        """
+        Convert the state message to the format required by the ConcurrentCursor.
+
+        e.g.
+        {
+            "state_type": ConcurrencyCompatibleStateType.date_range.value,
+            "metadata": { … },
+            "slices": [
+                {starts: 0, end: 1617030403, finished_processing: true}]
+        }
+        """
+        ...
+
+    @abstractmethod
+    def increment(self, value: Any) -> Any:
+        """
+        Increment a timestamp by a single unit.
+        """
+        ...
+
+    @abstractmethod
+    def output_format(self, value: Any) -> Any:
+        """
+        Convert the cursor value type to a JSON valid type.
+        """
+        ...
+
+    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:
+        """
+        Parse the value of the cursor field into a comparable value.
+        """
+        ...
+
+    @property
+    @abstractmethod
+    def zero_value(self) -> Any: ...