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

airbyte_cdk/sources/streams/checkpoint/checkpoint_reader.py

@@ -0,0 +1,335 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any, Iterable, Mapping, Optional
+
+from airbyte_cdk.sources.types import StreamSlice
+
+from .cursor import Cursor
+
+
+class CheckpointMode(Enum):
+    INCREMENTAL = "incremental"
+    RESUMABLE_FULL_REFRESH = "resumable_full_refresh"
+    FULL_REFRESH = "full_refresh"
+
+
+FULL_REFRESH_COMPLETE_STATE: Mapping[str, Any] = {"__ab_full_refresh_sync_complete": True}
+
+
+class CheckpointReader(ABC):
+    """
+    CheckpointReader manages how to iterate over a stream's partitions and serves as the bridge for interpreting the current state
+    of the stream that should be emitted back to the platform.
+    """
+
+    @abstractmethod
+    def next(self) -> Optional[Mapping[str, Any]]:
+        """
+        Returns the next slice that will be used to fetch the next group of records. Returning None indicates that the reader
+        has finished iterating over all slices.
+        """
+
+    @abstractmethod
+    def observe(self, new_state: Mapping[str, Any]) -> None:
+        """
+        Updates the internal state of the checkpoint reader based on the incoming stream state from a connector.
+
+        WARNING: This is used to retain backwards compatibility with streams using the legacy get_stream_state() method.
+        In order to uptake Resumable Full Refresh, connectors must migrate streams to use the state setter/getter methods.
+        """
+
+    @abstractmethod
+    def get_checkpoint(self) -> Optional[Mapping[str, Any]]:
+        """
+        Retrieves the current state value of the stream. The connector does not emit state messages if the checkpoint value is None.
+        """
+
+
+class IncrementalCheckpointReader(CheckpointReader):
+    """
+    IncrementalCheckpointReader handles iterating through a stream based on partitioned windows of data that are determined
+    before syncing data.
+    """
+
+    def __init__(
+        self, stream_state: Mapping[str, Any], stream_slices: Iterable[Optional[Mapping[str, Any]]]
+    ):
+        self._state: Optional[Mapping[str, Any]] = stream_state
+        self._stream_slices = iter(stream_slices)
+        self._has_slices = False
+
+    def next(self) -> Optional[Mapping[str, Any]]:
+        try:
+            next_slice = next(self._stream_slices)
+            self._has_slices = True
+            return next_slice
+        except StopIteration:
+            # This is used to avoid sending a duplicate state message at the end of a sync since the stream has already
+            # emitted state at the end of each slice. If we want to avoid this extra complexity, we can also just accept
+            # that every sync emits a final duplicate state
+            if self._has_slices:
+                self._state = None
+            return None
+
+    def observe(self, new_state: Mapping[str, Any]) -> None:
+        self._state = new_state
+
+    def get_checkpoint(self) -> Optional[Mapping[str, Any]]:
+        return self._state
+
+
+class CursorBasedCheckpointReader(CheckpointReader):
+    """
+    CursorBasedCheckpointReader is used by streams that implement a Cursor in order to manage state. This allows the checkpoint
+    reader to delegate the complexity of fetching state to the cursor and focus on the iteration over a stream's partitions.
+
+    This reader supports the Cursor interface used by Python and low-code sources. Not to be confused with Cursor interface
+    that belongs to the Concurrent CDK.
+    """
+
+    def __init__(
+        self,
+        cursor: Cursor,
+        stream_slices: Iterable[Optional[Mapping[str, Any]]],
+        read_state_from_cursor: bool = False,
+    ):
+        self._cursor = cursor
+        self._stream_slices = iter(stream_slices)
+        # read_state_from_cursor is used to delineate that partitions should determine when to stop syncing dynamically according
+        # to the value of the state at runtime. This currently only applies to streams that use resumable full refresh.
+        self._read_state_from_cursor = read_state_from_cursor
+        self._current_slice: Optional[StreamSlice] = None
+        self._finished_sync = False
+        self._previous_state: Optional[Mapping[str, Any]] = None
+
+    def next(self) -> Optional[Mapping[str, Any]]:
+        try:
+            self.current_slice = self._find_next_slice()
+            return self.current_slice
+        except StopIteration:
+            self._finished_sync = True
+            return None
+
+    def observe(self, new_state: Mapping[str, Any]) -> None:
+        # Cursor based checkpoint readers don't need to observe the new state because it has already been updated by the cursor
+        # while processing records
+        pass
+
+    def get_checkpoint(self) -> Optional[Mapping[str, Any]]:
+        # This is used to avoid sending a duplicate state messages
+        new_state = self._cursor.get_stream_state()
+        if new_state != self._previous_state:
+            self._previous_state = new_state
+            return new_state
+        else:
+            return None
+
+    def _find_next_slice(self) -> StreamSlice:
+        """
+        _find_next_slice() returns the next slice of data should be synced for the current stream according to its cursor.
+        This function supports iterating over a stream's slices across two dimensions. The first dimension is the stream's
+        partitions like parent records for a substream. The inner dimension iterates over the cursor value like a date
+        range for incremental streams or a pagination checkpoint for resumable full refresh.
+
+        The basic algorithm for iterating through a stream's slices is:
+        1. The first time next() is invoked we get the first partition
+        2. If the current partition is already complete as a result of a previous sync attempt, continue iterating until
+           we find an un-synced partition.
+        2. For streams whose cursor value is determined dynamically using stream state
+            1. Get the state for the current partition
+            2. If the current partition's state is complete, continue iterating over partitions
+            3. If the current partition's state is still in progress, emit the next cursor value
+            4. If the current partition is complete as delineated by the sentinel value, get the next incomplete partition
+        3. When stream has processed all partitions, the iterator will raise a StopIteration exception signaling there are no more
+           slices left for extracting more records.
+        """
+
+        if self._read_state_from_cursor:
+            if self.current_slice is None:
+                # current_slice is None represents the first time we are iterating over a stream's slices. The first slice to
+                # sync not been assigned yet and must first be read from the iterator
+                next_slice = self.read_and_convert_slice()
+                state_for_slice = self._cursor.select_state(next_slice)
+                if state_for_slice == FULL_REFRESH_COMPLETE_STATE:
+                    # Skip every slice that already has the terminal complete value indicating that a previous attempt
+                    # successfully synced the slice
+                    has_more = True
+                    while has_more:
+                        next_slice = self.read_and_convert_slice()
+                        state_for_slice = self._cursor.select_state(next_slice)
+                        has_more = state_for_slice == FULL_REFRESH_COMPLETE_STATE
+                return StreamSlice(
+                    cursor_slice=state_for_slice or {},
+                    partition=next_slice.partition,
+                    extra_fields=next_slice.extra_fields,
+                )
+            else:
+                state_for_slice = self._cursor.select_state(self.current_slice)
+                if state_for_slice == FULL_REFRESH_COMPLETE_STATE:
+                    # If the current slice is is complete, move to the next slice and skip the next slices that already
+                    # have the terminal complete value indicating that a previous attempt was successfully read.
+                    # Dummy initialization for mypy since we'll iterate at least once to get the next slice
+                    next_candidate_slice = StreamSlice(cursor_slice={}, partition={})
+                    has_more = True
+                    while has_more:
+                        next_candidate_slice = self.read_and_convert_slice()
+                        state_for_slice = self._cursor.select_state(next_candidate_slice)
+                        has_more = state_for_slice == FULL_REFRESH_COMPLETE_STATE
+                    return StreamSlice(
+                        cursor_slice=state_for_slice or {},
+                        partition=next_candidate_slice.partition,
+                        extra_fields=next_candidate_slice.extra_fields,
+                    )
+                # The reader continues to process the current partition if it's state is still in progress
+                return StreamSlice(
+                    cursor_slice=state_for_slice or {},
+                    partition=self.current_slice.partition,
+                    extra_fields=self.current_slice.extra_fields,
+                )
+        else:
+            # Unlike RFR cursors that iterate dynamically according to how stream state is updated, most cursors operate
+            # on a fixed set of slices determined before reading records. They just iterate to the next slice
+            return self.read_and_convert_slice()
+
+    @property
+    def current_slice(self) -> Optional[StreamSlice]:
+        return self._current_slice
+
+    @current_slice.setter
+    def current_slice(self, value: StreamSlice) -> None:
+        self._current_slice = value
+
+    def read_and_convert_slice(self) -> StreamSlice:
+        next_slice = next(self._stream_slices)
+        if not isinstance(next_slice, StreamSlice):
+            raise ValueError(
+                f"{self.current_slice} should be of type StreamSlice. This is likely a bug in the CDK, please contact Airbyte support"
+            )
+        return next_slice
+
+
+class LegacyCursorBasedCheckpointReader(CursorBasedCheckpointReader):
+    """
+    This (unfortunate) class operates like an adapter to retain backwards compatibility with legacy sources that take in stream_slice
+    in the form of a Mapping instead of the StreamSlice object. Internally, the reader still operates over StreamSlices, but it
+    is instantiated with and emits stream slices in the form of a Mapping[str, Any]. The logic of how partitions and cursors
+    are iterated over is synonymous with CursorBasedCheckpointReader.
+
+    We also retain the existing top level fields defined by the connector so the fields are present on dependent methods. For example,
+    the resulting mapping structure passed back to the stream's read_records() method looks like:
+    {
+      "cursor_slice": {
+        "next_page_token": 10
+      },
+      "partition": {
+        "repository": "airbytehq/airbyte"
+      },
+      "next_page_token": 10,
+      "repository": "airbytehq/airbyte"
+    }
+    """
+
+    def __init__(
+        self,
+        cursor: Cursor,
+        stream_slices: Iterable[Optional[Mapping[str, Any]]],
+        read_state_from_cursor: bool = False,
+    ):
+        super().__init__(
+            cursor=cursor,
+            stream_slices=stream_slices,
+            read_state_from_cursor=read_state_from_cursor,
+        )
+
+    def next(self) -> Optional[Mapping[str, Any]]:
+        try:
+            self.current_slice = self._find_next_slice()
+
+            if "partition" in dict(self.current_slice):
+                raise ValueError("Stream is configured to use invalid stream slice key 'partition'")
+            elif "cursor_slice" in dict(self.current_slice):
+                raise ValueError(
+                    "Stream is configured to use invalid stream slice key 'cursor_slice'"
+                )
+
+            # We convert StreamSlice to a regular mapping because legacy connectors operate on the basic Mapping object. We
+            # also duplicate all fields at the top level for backwards compatibility for existing Python sources
+            return {
+                "partition": self.current_slice.partition,
+                "cursor_slice": self.current_slice.cursor_slice,
+                **dict(self.current_slice),
+            }
+        except StopIteration:
+            self._finished_sync = True
+            return None
+
+    def read_and_convert_slice(self) -> StreamSlice:
+        next_mapping_slice = next(self._stream_slices)
+        if not isinstance(next_mapping_slice, Mapping):
+            raise ValueError(
+                f"{self.current_slice} should be of type Mapping. This is likely a bug in the CDK, please contact Airbyte support"
+            )
+
+        # The legacy reader is instantiated with an iterable of stream slice mappings. We convert each into a StreamSlice
+        # to sanely process them during the sync and to reuse the existing Python defined cursors
+        return StreamSlice(
+            partition=next_mapping_slice,
+            cursor_slice={},
+        )
+
+
+class ResumableFullRefreshCheckpointReader(CheckpointReader):
+    """
+    ResumableFullRefreshCheckpointReader allows for iteration over an unbounded set of records based on the pagination strategy
+    of the stream. Because the number of pages is unknown, the stream's current state is used to determine whether to continue
+    fetching more pages or stopping the sync.
+    """
+
+    def __init__(self, stream_state: Mapping[str, Any]):
+        # The first attempt of an RFR stream has an empty {} incoming state, but should still make a first attempt to read records
+        # from the first page in next().
+        self._first_page = bool(stream_state == {})
+        self._state: Mapping[str, Any] = stream_state
+
+    def next(self) -> Optional[Mapping[str, Any]]:
+        if self._first_page:
+            self._first_page = False
+            return self._state
+        elif self._state == FULL_REFRESH_COMPLETE_STATE:
+            return None
+        else:
+            return self._state
+
+    def observe(self, new_state: Mapping[str, Any]) -> None:
+        self._state = new_state
+
+    def get_checkpoint(self) -> Optional[Mapping[str, Any]]:
+        return self._state or {}
+
+
+class FullRefreshCheckpointReader(CheckpointReader):
+    """
+    FullRefreshCheckpointReader iterates over data that cannot be checkpointed incrementally during the sync because the stream
+    is not capable of managing state. At the end of a sync, a final state message is emitted to signal completion.
+    """
+
+    def __init__(self, stream_slices: Iterable[Optional[Mapping[str, Any]]]):
+        self._stream_slices = iter(stream_slices)
+        self._final_checkpoint = False
+
+    def next(self) -> Optional[Mapping[str, Any]]:
+        try:
+            return next(self._stream_slices)
+        except StopIteration:
+            self._final_checkpoint = True
+            return None
+
+    def observe(self, new_state: Mapping[str, Any]) -> None:
+        pass
+
+    def get_checkpoint(self) -> Optional[Mapping[str, Any]]:
+        if self._final_checkpoint:
+            return {"__ab_no_cursor_state_message": True}
+        return None

airbyte_cdk/sources/{declarative/incremental → streams/checkpoint}/cursor.py

@@ -3,16 +3,15 @@
 #
 
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Any, Optional
 
-from airbyte_cdk.sources.declarative.stream_slicers.stream_slicer import StreamSlicer
-from airbyte_cdk.sources.declarative.types import Record, StreamSlice, StreamState
+from airbyte_cdk.sources.types import Record, StreamSlice, StreamState
 
 
-class Cursor(ABC, StreamSlicer):
+class Cursor(ABC):
     """
-    Cursors are components that allow for incremental syncs. They keep track of what data has been consumed and slices the requests based on
-    that information.
+    Cursors are components that allow for checkpointing the current state of a sync. They keep track of what data has been consumed
+    and allows for syncs to be resumed from a specific point based on that information.
     """
 
     @abstractmethod
@@ -35,17 +34,13 @@ class Cursor(ABC, StreamSlicer):
         pass
 
     @abstractmethod
-    def close_slice(self, stream_slice: StreamSlice, most_recent_record: Optional[Record]) -> None:
+    def close_slice(self, stream_slice: StreamSlice, *args: Any) -> None:
         """
-        Update state based on the stream slice and the latest record. Note that `stream_slice.cursor_slice` and
-        `most_recent_record.associated_slice` are expected to be the same but we make it explicit here that `stream_slice` should be leveraged to
-        update the state.
+        Update state based on the stream slice. Note that `stream_slice.cursor_slice` and `most_recent_record.associated_slice` are expected
+        to be the same but we make it explicit here that `stream_slice` should be leveraged to update the state. We do not pass in the
+        latest record, since cursor instances should maintain the relevant internal state on their own.
 
         :param stream_slice: slice to close
-        :param most_recent_record: the latest record we have received for the slice. This is important to consider because even if the
-          cursor emits a slice, some APIs are not able to enforce the upper boundary. The outcome is that the last_record might have a
-          higher cursor value than the slice upper boundary and if we want to reduce the duplication as much as possible, we need to
-          consider the highest value between the internal cursor, the stream slice upper boundary and the record cursor value.
         """
 
     @abstractmethod
@@ -72,3 +67,11 @@ class Cursor(ABC, StreamSlicer):
         """
         Evaluating which record is greater in terms of cursor. This is used to avoid having to capture all the records to close a slice
         """
+
+    @abstractmethod
+    def select_state(self, stream_slice: Optional[StreamSlice] = None) -> Optional[StreamState]:
+        """
+        Get the state value of a specific stream_slice. For incremental or resumable full refresh cursors which only manage state in
+        a single dimension this is the entire state object. For per-partition cursors used by substreams, this returns the state of
+        a specific parent delineated by the incoming slice's partition object.
+        """

airbyte_cdk/sources/streams/checkpoint/per_partition_key_serializer.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+import json
+from typing import Any, Mapping
+
+
+class PerPartitionKeySerializer:
+    """
+    We are concerned of the performance of looping through the `states` list and evaluating equality on the partition. To reduce this
+    concern, we wanted to use dictionaries to map `partition -> cursor`. However, partitions are dict and dict can't be used as dict keys
+    since they are not hashable. By creating json string using the dict, we can have a use the dict as a key to the dict since strings are
+    hashable.
+    """
+
+    @staticmethod
+    def to_partition_key(to_serialize: Any) -> str:
+        # separators have changed in Python 3.4. To avoid being impacted by further change, we explicitly specify our own value
+        return json.dumps(to_serialize, indent=None, separators=(",", ":"), sort_keys=True)
+
+    @staticmethod
+    def to_partition(to_deserialize: Any) -> Mapping[str, Any]:
+        return json.loads(to_deserialize)  # type: ignore # The partition is known to be a dict, but the type hint is Any

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.