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

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()