airbyte-cdk 0.72.0__py3-none-any.whl → 6.13.1.dev4106__py3-none-any.whl

airbyte_cdk/sources/streams/checkpoint/resumable_full_refresh_cursor.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from airbyte_cdk.sources.streams.checkpoint import Cursor
+from airbyte_cdk.sources.types import Record, StreamSlice, StreamState
+
+
+@dataclass
+class ResumableFullRefreshCursor(Cursor):
+    """
+    Cursor that allows for the checkpointing of sync progress according to a synthetic cursor based on the pagination state
+    of the stream. Resumable full refresh syncs are only intended to retain state in between sync attempts of the same job
+    with the platform responsible for removing said state.
+    """
+
+    def __init__(self) -> None:
+        self._cursor: StreamState = {}
+
+    def get_stream_state(self) -> StreamState:
+        return self._cursor
+
+    def set_initial_state(self, stream_state: StreamState) -> None:
+        self._cursor = stream_state
+
+    def observe(self, stream_slice: StreamSlice, record: Record) -> None:
+        """
+        Resumable full refresh manages state using a page number so it does not need to update state by observing incoming records.
+        """
+        pass
+
+    def close_slice(self, stream_slice: StreamSlice, *args: Any) -> None:
+        self._cursor = stream_slice.cursor_slice
+
+    def should_be_synced(self, record: Record) -> bool:
+        """
+        Unlike date-based cursors which filter out records outside slice boundaries, resumable full refresh records exist within pages
+        that don't have filterable bounds. We should always return them.
+        """
+        return True
+
+    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
+        """
+        RFR record don't have ordering to be compared between one another.
+        """
+        return False
+
+    def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
+        # A top-level RFR cursor only manages the state of a single partition
+        return self._cursor

airbyte_cdk/sources/streams/checkpoint/substream_resumable_full_refresh_cursor.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from dataclasses import dataclass
+from typing import Any, Mapping, MutableMapping, Optional
+
+from airbyte_cdk.models import FailureType
+from airbyte_cdk.sources.streams.checkpoint import Cursor
+from airbyte_cdk.sources.streams.checkpoint.per_partition_key_serializer import (
+    PerPartitionKeySerializer,
+)
+from airbyte_cdk.sources.types import Record, StreamSlice, StreamState
+from airbyte_cdk.utils import AirbyteTracedException
+
+FULL_REFRESH_COMPLETE_STATE: Mapping[str, Any] = {"__ab_full_refresh_sync_complete": True}
+
+
+@dataclass
+class SubstreamResumableFullRefreshCursor(Cursor):
+    def __init__(self) -> None:
+        self._per_partition_state: MutableMapping[str, StreamState] = {}
+        self._partition_serializer = PerPartitionKeySerializer()
+
+    def get_stream_state(self) -> StreamState:
+        return {"states": list(self._per_partition_state.values())}
+
+    def set_initial_state(self, stream_state: StreamState) -> None:
+        """
+        Set the initial state for the cursors.
+
+        This method initializes the state for each partition cursor using the provided stream state.
+        If a partition state is provided in the stream state, it will update the corresponding partition cursor with this state.
+
+        To simplify processing and state management, we do not maintain the checkpointed state of the parent partitions.
+        Instead, we are tracking whether a parent has already successfully synced on a prior attempt and skipping over it
+        allowing the sync to continue making progress. And this works for RFR because the platform will dispose of this
+        state on the next sync job.
+
+        Args:
+            stream_state (StreamState): The state of the streams to be set. The format of the stream state should be:
+                {
+                    "states": [
+                        {
+                            "partition": {
+                                "partition_key": "value_0"
+                            },
+                            "cursor": {
+                                "__ab_full_refresh_sync_complete": True
+                            }
+                        },
+                        {
+                            "partition": {
+                                "partition_key": "value_1"
+                            },
+                            "cursor": {},
+                        },
+                    ]
+                }
+        """
+        if not stream_state:
+            return
+
+        if "states" not in stream_state:
+            raise AirbyteTracedException(
+                internal_message=f"Could not sync parse the following state: {stream_state}",
+                message="The state for is format invalid. Validate that the migration steps included a reset and that it was performed "
+                "properly. Otherwise, please contact Airbyte support.",
+                failure_type=FailureType.config_error,
+            )
+
+        for state in stream_state["states"]:
+            self._per_partition_state[self._to_partition_key(state["partition"])] = state
+
+    def observe(self, stream_slice: StreamSlice, record: Record) -> None:
+        """
+        Substream resumable full refresh manages state by closing the slice after syncing a parent so observe is not used.
+        """
+        pass
+
+    def close_slice(self, stream_slice: StreamSlice, *args: Any) -> None:
+        self._per_partition_state[self._to_partition_key(stream_slice.partition)] = {
+            "partition": stream_slice.partition,
+            "cursor": FULL_REFRESH_COMPLETE_STATE,
+        }
+
+    def should_be_synced(self, record: Record) -> bool:
+        """
+        Unlike date-based cursors which filter out records outside slice boundaries, resumable full refresh records exist within pages
+        that don't have filterable bounds. We should always return them.
+        """
+        return True
+
+    def is_greater_than_or_equal(self, first: Record, second: Record) -> bool:
+        """
+        RFR record don't have ordering to be compared between one another.
+        """
+        return False
+
+    def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
+        if not stream_slice:
+            raise ValueError("A partition needs to be provided in order to extract a state")
+
+        return self._per_partition_state.get(
+            self._to_partition_key(stream_slice.partition), {}
+        ).get("cursor")
+
+    def _to_partition_key(self, partition: Mapping[str, Any]) -> str:
+        return self._partition_serializer.to_partition_key(partition)
+
+    def _to_dict(self, partition_key: str) -> Mapping[str, Any]:
+        return self._partition_serializer.to_partition(partition_key)

airbyte_cdk/sources/streams/concurrent/README.md

@@ -0,0 +1,7 @@
+## Breaking Changes & Limitations
+
+- [bigger scope than Concurrent CDK] checkpointing state was acting on the number of records per slice. This has been changed to consider the number of records per syncs
+- `Source.read_state` and `Source._emit_legacy_state_format` are now classmethods to allow for developers to have access to the state before instantiating the source
+- send_per_stream_state is always True for Concurrent CDK
+- Using stream_state during read_records: The concern is that today, stream_instance.get_updated_state is called on every record and read_records on every slice. The implication is that the argument stream_state passed to read_records will have the value after the last stream_instance.get_updated_state of the previous slice. For Concurrent CDK, this is not possible as slices are processed in an unordered way.
+- Cursor fields can only be data-time formatted as epoch. Eventually, we want to move to ISO 8601 as it provides more flexibility but for the first iteration on Stripe, it was easier to use the same format that was already used

airbyte_cdk/sources/streams/concurrent/abstract_stream.py

@@ -5,14 +5,19 @@
 from abc import ABC, abstractmethod
 from typing import Any, Iterable, Mapping, Optional
 
+from typing_extensions import deprecated
+
 from airbyte_cdk.models import AirbyteStream
+from airbyte_cdk.sources.source import ExperimentalClassWarning
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import StreamAvailability
 from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from deprecated.classic import deprecated
 
 
-@deprecated("This class is experimental. Use at your own risk.")
+@deprecated(
+    "This class is experimental. Use at your own risk.",
+    category=ExperimentalClassWarning,
+)
 class AbstractStream(ABC):
     """
     AbstractStream is an experimental interface for streams developed as part of the Concurrent CDK.

airbyte_cdk/sources/streams/concurrent/adapters.py

@@ -8,37 +8,52 @@ import logging
 from functools import lru_cache
 from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union
 
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteStream, ConfiguredAirbyteStream, Level, SyncMode, Type
+from typing_extensions import deprecated
+
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteStream,
+    ConfiguredAirbyteStream,
+    Level,
+    SyncMode,
+    Type,
+)
 from airbyte_cdk.sources import AbstractSource, Source
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.source import ExperimentalClassWarning
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.availability_strategy import AvailabilityStrategy
 from airbyte_cdk.sources.streams.concurrent.abstract_stream_facade import AbstractStreamFacade
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import (
     AbstractAvailabilityStrategy,
-    StreamAvailability,
-    StreamAvailable,
-    StreamUnavailable,
+    AlwaysAvailableAvailabilityStrategy,
 )
 from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, FinalStateCursor
 from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 from airbyte_cdk.sources.streams.concurrent.exceptions import ExceptionWithDisplayMessage
-from airbyte_cdk.sources.streams.concurrent.helpers import get_cursor_field_from_stream, get_primary_key_from_stream
+from airbyte_cdk.sources.streams.concurrent.helpers import (
+    get_cursor_field_from_stream,
+    get_primary_key_from_stream,
+)
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 from airbyte_cdk.sources.streams.core import StreamData
+from airbyte_cdk.sources.types import Record
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig
 from airbyte_cdk.sources.utils.slice_logger import SliceLogger
-from deprecated.classic import deprecated
+from airbyte_cdk.utils.slice_hasher import SliceHasher
 
 """
 This module contains adapters to help enabling concurrency on Stream objects without needing to migrate to AbstractStream
 """
 
 
-@deprecated("This class is experimental. Use at your own risk.")
+@deprecated(
+    "This class is experimental. Use at your own risk.",
+    category=ExperimentalClassWarning,
+)
 class StreamFacade(AbstractStreamFacade[DefaultStream], Stream):
     """
     The StreamFacade is a Stream that wraps an AbstractStream and exposes it as a Stream.
@@ -77,15 +92,16 @@ class StreamFacade(AbstractStreamFacade[DefaultStream], Stream):
                 partition_generator=StreamPartitionGenerator(
                     stream,
                     message_repository,
-                    SyncMode.full_refresh if isinstance(cursor, FinalStateCursor) else SyncMode.incremental,
+                    SyncMode.full_refresh
+                    if isinstance(cursor, FinalStateCursor)
+                    else SyncMode.incremental,
                     [cursor_field] if cursor_field is not None else None,
                     state,
-                    cursor,
                 ),
                 name=stream.name,
                 namespace=stream.namespace,
                 json_schema=stream.get_json_schema(),
-                availability_strategy=StreamAvailabilityStrategy(stream, source),
+                availability_strategy=AlwaysAvailableAvailabilityStrategy(),
                 primary_key=pk,
                 cursor_field=cursor_field,
                 logger=logger,
@@ -99,14 +115,23 @@ class StreamFacade(AbstractStreamFacade[DefaultStream], Stream):
 
     @property
     def state(self) -> MutableMapping[str, Any]:
-        raise NotImplementedError("This should not be called as part of the Concurrent CDK code. Please report the problem to Airbyte")
+        raise NotImplementedError(
+            "This should not be called as part of the Concurrent CDK code. Please report the problem to Airbyte"
+        )
 
     @state.setter
     def state(self, value: Mapping[str, Any]) -> None:
         if "state" in dir(self._legacy_stream):
             self._legacy_stream.state = value  # type: ignore  # validating `state` is attribute of stream using `if` above
 
-    def __init__(self, stream: DefaultStream, legacy_stream: Stream, cursor: Cursor, slice_logger: SliceLogger, logger: logging.Logger):
+    def __init__(
+        self,
+        stream: DefaultStream,
+        legacy_stream: Stream,
+        cursor: Cursor,
+        slice_logger: SliceLogger,
+        logger: logging.Logger,
+    ):
         """
         :param stream: The underlying AbstractStream
         """
@@ -143,7 +168,10 @@ class StreamFacade(AbstractStreamFacade[DefaultStream], Stream):
                 # This shouldn't happen if the ConcurrentCursor was used
                 state = "unknown; no state attribute was available on the cursor"
             yield AirbyteMessage(
-                type=Type.LOG, log=AirbyteLogMessage(level=Level.ERROR, message=f"Cursor State at time of exception: {state}")
+                type=Type.LOG,
+                log=AirbyteLogMessage(
+                    level=Level.ERROR, message=f"Cursor State at time of exception: {state}"
+                ),
             )
             raise exc
 
@@ -170,6 +198,10 @@ class StreamFacade(AbstractStreamFacade[DefaultStream], Stream):
         else:
             return self._abstract_stream.cursor_field
 
+    @property
+    def cursor(self) -> Optional[Cursor]:  # type: ignore[override] # StreamFaced expects to use only airbyte_cdk.sources.streams.concurrent.cursor.Cursor
+        return self._cursor
+
     @lru_cache(maxsize=None)
     def get_json_schema(self) -> Mapping[str, Any]:
         return self._abstract_stream.get_json_schema()
@@ -178,7 +210,9 @@ class StreamFacade(AbstractStreamFacade[DefaultStream], Stream):
     def supports_incremental(self) -> bool:
         return self._legacy_stream.supports_incremental
 
-    def check_availability(self, logger: logging.Logger, source: Optional["Source"] = None) -> Tuple[bool, Optional[str]]:
+    def check_availability(
+        self, logger: logging.Logger, source: Optional["Source"] = None
+    ) -> Tuple[bool, Optional[str]]:
         """
         Verifies the stream is available. Delegates to the underlying AbstractStream and ignores the parameters
         :param logger: (ignored)
@@ -198,6 +232,15 @@ class StreamFacade(AbstractStreamFacade[DefaultStream], Stream):
         return self._abstract_stream
 
 
+class SliceEncoder(json.JSONEncoder):
+    def default(self, obj: Any) -> Any:
+        if hasattr(obj, "__json_serializable__"):
+            return obj.__json_serializable__()
+
+        # Let the base class default method raise the TypeError
+        return super().default(obj)
+
+
 class StreamPartition(Partition):
     """
     This class acts as an adapter between the new Partition interface and the Stream's stream_slice interface
@@ -216,7 +259,6 @@ class StreamPartition(Partition):
         sync_mode: SyncMode,
         cursor_field: Optional[List[str]],
         state: Optional[MutableMapping[str, Any]],
-        cursor: Cursor,
     ):
         """
         :param stream: The stream to delegate to
@@ -229,8 +271,7 @@ class StreamPartition(Partition):
         self._sync_mode = sync_mode
         self._cursor_field = cursor_field
         self._state = state
-        self._cursor = cursor
-        self._is_closed = False
+        self._hash = SliceHasher.hash(self._stream.name, self._slice)
 
     def read(self) -> Iterable[Record]:
         """
@@ -253,10 +294,14 @@ class StreamPartition(Partition):
             ):
                 if isinstance(record_data, Mapping):
                     data_to_return = dict(record_data)
-                    self._stream.transformer.transform(data_to_return, self._stream.get_json_schema())
-                    record = Record(data_to_return, self._stream.name)
-                    self._cursor.observe(record)
-                    yield Record(data_to_return, self._stream.name)
+                    self._stream.transformer.transform(
+                        data_to_return, self._stream.get_json_schema()
+                    )
+                    yield Record(
+                        data=data_to_return,
+                        stream_name=self.stream_name(),
+                        associated_slice=self._slice,  # type: ignore [arg-type]
+                    )
                 else:
                     self._message_repository.emit_message(record_data)
         except Exception as e:
@@ -270,23 +315,11 @@ class StreamPartition(Partition):
         return self._slice
 
     def __hash__(self) -> int:
-        if self._slice:
-            # Convert the slice to a string so that it can be hashed
-            s = json.dumps(self._slice, sort_keys=True)
-            return hash((self._stream.name, s))
-        else:
-            return hash(self._stream.name)
+        return self._hash
 
     def stream_name(self) -> str:
         return self._stream.name
 
-    def close(self) -> None:
-        self._cursor.close_partition(self)
-        self._is_closed = True
-
-    def is_closed(self) -> bool:
-        return self._is_closed
-
     def __repr__(self) -> str:
         return f"StreamPartition({self._stream.name}, {self._slice})"
 
@@ -306,7 +339,6 @@ class StreamPartitionGenerator(PartitionGenerator):
         sync_mode: SyncMode,
         cursor_field: Optional[List[str]],
         state: Optional[MutableMapping[str, Any]],
-        cursor: Cursor,
     ):
         """
         :param stream: The stream to delegate to
@@ -317,21 +349,32 @@ class StreamPartitionGenerator(PartitionGenerator):
         self._sync_mode = sync_mode
         self._cursor_field = cursor_field
         self._state = state
-        self._cursor = cursor
 
     def generate(self) -> Iterable[Partition]:
-        for s in self._stream.stream_slices(sync_mode=self._sync_mode, cursor_field=self._cursor_field, stream_state=self._state):
+        for s in self._stream.stream_slices(
+            sync_mode=self._sync_mode, cursor_field=self._cursor_field, stream_state=self._state
+        ):
             yield StreamPartition(
-                self._stream, copy.deepcopy(s), self.message_repository, self._sync_mode, self._cursor_field, self._state, self._cursor
+                self._stream,
+                copy.deepcopy(s),
+                self.message_repository,
+                self._sync_mode,
+                self._cursor_field,
+                self._state,
             )
 
 
-@deprecated("This class is experimental. Use at your own risk.")
+@deprecated(
+    "Availability strategy has been soft deprecated. Do not use. Class is subject to removal",
+    category=ExperimentalClassWarning,
+)
 class AvailabilityStrategyFacade(AvailabilityStrategy):
     def __init__(self, abstract_availability_strategy: AbstractAvailabilityStrategy):
         self._abstract_availability_strategy = abstract_availability_strategy
 
-    def check_availability(self, stream: Stream, logger: logging.Logger, source: Optional[Source]) -> Tuple[bool, Optional[str]]:
+    def check_availability(
+        self, stream: Stream, logger: logging.Logger, source: Optional["Source"] = None
+    ) -> Tuple[bool, Optional[str]]:
         """
         Checks stream availability.
 
@@ -344,37 +387,3 @@ class AvailabilityStrategyFacade(AvailabilityStrategy):
         """
         stream_availability = self._abstract_availability_strategy.check_availability(logger)
         return stream_availability.is_available(), stream_availability.message()
-
-
-class StreamAvailabilityStrategy(AbstractAvailabilityStrategy):
-    """
-    This class acts as an adapter between the existing AvailabilityStrategy and the new AbstractAvailabilityStrategy.
-    StreamAvailabilityStrategy is instantiated with a Stream and a Source to allow the existing AvailabilityStrategy to be used with the new AbstractAvailabilityStrategy interface.
-
-    A more convenient implementation would not depend on the docs URL instead of the Source itself, and would support running on an AbstractStream instead of only on a Stream.
-
-    This class can be used to help enable concurrency on existing connectors without having to rewrite everything as AbstractStream and AbstractAvailabilityStrategy.
-    In the long-run, it would be preferable to update the connectors, but we don't have the tooling or need to justify the effort at this time.
-    """
-
-    def __init__(self, stream: Stream, source: Source):
-        """
-        :param stream: The stream to delegate to
-        :param source: The source to delegate to
-        """
-        self._stream = stream
-        self._source = source
-
-    def check_availability(self, logger: logging.Logger) -> StreamAvailability:
-        try:
-            available, message = self._stream.check_availability(logger, self._source)
-            if available:
-                return StreamAvailable()
-            else:
-                return StreamUnavailable(str(message))
-        except Exception as e:
-            display_message = self._stream.get_error_display_message(e)
-            if display_message:
-                raise ExceptionWithDisplayMessage(display_message)
-            else:
-                raise e

airbyte_cdk/sources/streams/concurrent/availability_strategy.py

@@ -6,7 +6,9 @@ import logging
 from abc import ABC, abstractmethod
 from typing import Optional
 
-from deprecated.classic import deprecated
+from typing_extensions import deprecated
+
+from airbyte_cdk.sources.source import ExperimentalClassWarning
 
 
 class StreamAvailability(ABC):
@@ -46,7 +48,10 @@ class StreamUnavailable(StreamAvailability):
 STREAM_AVAILABLE = StreamAvailable()
 
 
-@deprecated("This class is experimental. Use at your own risk.")
+@deprecated(
+    "This class is experimental. Use at your own risk.",
+    category=ExperimentalClassWarning,
+)
 class AbstractAvailabilityStrategy(ABC):
     """
     AbstractAvailabilityStrategy is an experimental interface developed as part of the Concurrent CDK.
@@ -64,3 +69,26 @@ class AbstractAvailabilityStrategy(ABC):
         :param logger: logger object to use
         :return: A StreamAvailability object describing the stream's availability
         """
+
+
+@deprecated(
+    "This class is experimental. Use at your own risk.",
+    category=ExperimentalClassWarning,
+)
+class AlwaysAvailableAvailabilityStrategy(AbstractAvailabilityStrategy):
+    """
+    An availability strategy that always indicates a stream is available.
+
+    This strategy is used to avoid breaking changes and serves as a soft
+    deprecation of the availability strategy, allowing a smoother transition
+    without disrupting existing functionality.
+    """
+
+    def check_availability(self, logger: logging.Logger) -> StreamAvailability:
+        """
+        Checks stream availability.
+
+        :param logger: logger object to use
+        :return: A StreamAvailability object describing the stream's availability
+        """
+        return StreamAvailable()