have declarative availability check support AbstractStream

Author: maxi297

airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py

@@ -3,12 +3,13 @@
 #
 
 import logging
-import traceback
 from dataclasses import InitVar, dataclass
-from typing import Any, List, Mapping, Tuple
+from typing import Any, List, Mapping, Tuple, Union
 
-from airbyte_cdk import AbstractSource
+from airbyte_cdk.sources.abstract_source import AbstractSource
+from airbyte_cdk.sources.declarative.checks.check_stream import evaluate_availability
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
+from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.http.availability_strategy import HttpAvailabilityStrategy
 
 
@@ -34,20 +35,16 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
     def check_connection(
         self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
     ) -> Tuple[bool, Any]:
-        streams = source.streams(config=config)
+        streams: List[Union[Stream, AbstractStream]] = source.streams(config=config)  # type: ignore  # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker
 
         if len(streams) == 0:
             return False, f"No streams to connect to from source {source}"
         if not self.use_check_availability:
             return True, None
 
-        availability_strategy = HttpAvailabilityStrategy()
-
         try:
             for stream in streams[: min(self.stream_count, len(streams))]:
-                stream_is_available, reason = availability_strategy.check_availability(
-                    stream, logger
-                )
+                stream_is_available, reason = evaluate_availability(stream, logger)
                 if not stream_is_available:
                     logger.warning(f"Stream {stream.name} is not available: {reason}")
                     return False, reason

airbyte_cdk/sources/declarative/checks/check_stream.py

@@ -5,13 +5,28 @@
 import logging
 import traceback
 from dataclasses import InitVar, dataclass
-from typing import Any, Dict, List, Mapping, Optional, Tuple
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
 
-from airbyte_cdk import AbstractSource
+from airbyte_cdk.sources.streams.core import Stream
+from airbyte_cdk.sources.abstract_source import AbstractSource
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
+from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.http.availability_strategy import HttpAvailabilityStrategy
 
 
+def evaluate_availability(stream: Union[Stream, AbstractStream], logger: logging.Logger) -> Tuple[bool, Optional[str]]:
+    """
+    As a transition period, we want to support both Stream and AbstractStream until we migrate everything to AbstractStream.
+    """
+    if isinstance(stream, Stream):
+        return HttpAvailabilityStrategy().check_availability(stream, logger)
+    elif isinstance(stream, AbstractStream):
+        availability = stream.check_availability()
+        return availability.is_available, availability.reason
+    else:
+        raise ValueError(f"Unsupported stream type {type(stream)}")
+
+
 @dataclass(frozen=True)
 class DynamicStreamCheckConfig:
     """Defines the configuration for dynamic stream during connection checking. This class specifies
@@ -51,7 +66,7 @@ def check_connection(
     ) -> Tuple[bool, Any]:
         """Checks the connection to the source and its streams."""
         try:
-            streams = source.streams(config=config)
+            streams: List[Union[Stream, AbstractStream]] = source.streams(config=config)  # type: ignore  # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker
             if not streams:
                 return False, f"No streams to connect to from source {source}"
         except Exception as error:
@@ -82,13 +97,12 @@ def check_connection(
         return True, None
 
     def _check_stream_availability(
-        self, stream_name_to_stream: Dict[str, Any], stream_name: str, logger: logging.Logger
+        self, stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]], stream_name: str, logger: logging.Logger
     ) -> Tuple[bool, Any]:
         """Checks if streams are available."""
-        availability_strategy = HttpAvailabilityStrategy()
         try:
             stream = stream_name_to_stream[stream_name]
-            stream_is_available, reason = availability_strategy.check_availability(stream, logger)
+            stream_is_available, reason = evaluate_availability(stream, logger)
             if not stream_is_available:
                 message = f"Stream {stream_name} is not available: {reason}"
                 logger.warning(message)
@@ -98,7 +112,7 @@ def _check_stream_availability(
         return True, None
 
     def _check_dynamic_streams_availability(
-        self, source: AbstractSource, stream_name_to_stream: Dict[str, Any], logger: logging.Logger
+        self, source: AbstractSource, stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]], logger: logging.Logger
     ) -> Tuple[bool, Any]:
         """Checks the availability of dynamic streams."""
         dynamic_streams = source.resolved_manifest.get("dynamic_streams", [])  # type: ignore[attr-defined] # The source's resolved_manifest manifest is checked before calling this method
@@ -135,18 +149,15 @@ def _map_generated_streams(
     def _check_generated_streams_availability(
         self,
         generated_streams: List[Dict[str, Any]],
-        stream_name_to_stream: Dict[str, Any],
+        stream_name_to_stream: Dict[str, Union[Stream, AbstractStream]],
         logger: logging.Logger,
         max_count: int,
     ) -> Tuple[bool, Any]:
         """Checks availability of generated dynamic streams."""
-        availability_strategy = HttpAvailabilityStrategy()
         for declarative_stream in generated_streams[: min(max_count, len(generated_streams))]:
             stream = stream_name_to_stream[declarative_stream["name"]]
             try:
-                stream_is_available, reason = availability_strategy.check_availability(
-                    stream, logger
-                )
+                stream_is_available, reason = evaluate_availability(stream, logger)
                 if not stream_is_available:
                     message = f"Dynamic Stream {stream.name} is not available: {reason}"
                     logger.warning(message)

airbyte_cdk/sources/streams/concurrent/abstract_stream.py

@@ -9,6 +9,7 @@
 
 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
 
@@ -87,3 +88,9 @@ def cursor(self) -> Cursor:
         """
         :return: The cursor associated with this stream.
         """
+
+    @abstractmethod
+    def check_availability(self) -> StreamAvailability:
+        """
+        :return: If the stream is available and if not, why
+        """

airbyte_cdk/sources/streams/concurrent/availability_strategy.py

@@ -0,0 +1,38 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+
+class StreamAvailability:
+
+    @classmethod
+    def available(cls) -> "StreamAvailability":
+        return cls(True)
+
+    @classmethod
+    def unavailable(cls, reason: str) -> "StreamAvailability":
+        return StreamAvailability(False, reason)
+
+    def __init__(self, available: bool, reason: Optional[str] = None) -> None:
+        self._available = available
+        self._reason = reason
+
+        if not available:
+            assert reason, "A reason needs to be provided if the stream is not available"
+
+    @property
+    def is_available(self) -> bool:
+        """
+        :return: True if the stream is available. False if the stream is not
+        """
+        return self._available
+
+    @property
+    def reason(self) -> Optional[str]:
+        """
+        :return: A message describing why the stream is not available. If the stream is available, this should return None.
+        """
+        return self._reason

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -8,12 +8,15 @@
 
 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 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
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
 class DefaultStream(AbstractStream):
+
     def __init__(
         self,
         partition_generator: PartitionGenerator,
@@ -91,3 +94,29 @@ def log_stream_sync_configuration(self) -> None:
     @property
     def cursor(self) -> Cursor:
         return self._cursor
+
+    def check_availability(self) -> StreamAvailability:
+        """
+        Check stream availability by attempting to read the first record of the stream.
+        """
+        try:
+            partition = next(iter(self.generate_partitions()))
+        except StopIteration:
+            # NOTE: The following comment was copied from legacy stuff and I don't know how relevant it is:
+            # If stream_slices has no `next()` item (Note - this is different from stream_slices returning [None]!)
+            # This can happen when a substream's `stream_slices` method does a `for record in parent_records: yield <something>`
+            # without accounting for the case in which the parent stream is empty.
+            return StreamAvailability.unavailable(
+                f"Cannot attempt to connect to stream {self.name} - no stream slices were found"
+            )
+        except AirbyteTracedException as error:
+            return StreamAvailability.unavailable(error.message)
+
+        try:
+            next(iter(partition.read()))
+            return StreamAvailability.available()
+        except StopIteration:
+            self._logger.info(f"Successfully connected to stream {self.name}, but got 0 records.")
+            return StreamAvailability.available()
+        except AirbyteTracedException as error:
+            return StreamAvailability.unavailable(error.message)

unit_tests/sources/declarative/checks/test_check_stream.py

@@ -17,6 +17,7 @@
     ConcurrentDeclarativeSource,
 )
 from airbyte_cdk.sources.streams.http import HttpStream
+from airbyte_cdk.sources.streams.core import Stream
 from airbyte_cdk.test.mock_http import HttpMocker, HttpRequest, HttpResponse
 
 logger = logging.getLogger("test")
@@ -45,7 +46,7 @@
 def test_check_stream_with_slices_as_list(
     test_name, record, streams_to_check, stream_slice, expectation, slices_as_list
 ):
-    stream = MagicMock()
+    stream = MagicMock(spec=Stream)
     stream.name = "s1"
     stream.availability_strategy = None
     if slices_as_list:
@@ -77,7 +78,7 @@ def mock_read_records(responses, default_response=None, **kwargs):
 
 
 def test_check_empty_stream():
-    stream = MagicMock()
+    stream = MagicMock(spec=Stream)
     stream.name = "s1"
     stream.read_records.return_value = iter([])
     stream.stream_slices.return_value = iter([None])
@@ -91,7 +92,7 @@ def test_check_empty_stream():
 
 
 def test_check_stream_with_no_stream_slices_aborts():
-    stream = MagicMock()
+    stream = MagicMock(spec=Stream)
     stream.name = "s1"
     stream.stream_slices.return_value = iter([])
 

unit_tests/sources/streams/concurrent/test_default_stream.py

@@ -4,15 +4,22 @@
 import unittest
 from unittest.mock import Mock
 
+import pytest
+
 from airbyte_cdk.models import AirbyteStream, SyncMode
 from airbyte_cdk.sources.message import InMemoryMessageRepository
 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.partitions.partition import Partition
+from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
+from airbyte_cdk.sources.types import Record
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
 class ThreadBasedConcurrentStreamTest(unittest.TestCase):
     def setUp(self):
-        self._partition_generator = Mock()
+        self._partition_generator = Mock(spec=PartitionGenerator)
+        self._partition = Mock(spec=Partition)
         self._name = "name"
         self._json_schema = {}
         self._primary_key = []
@@ -243,3 +250,62 @@ def test_as_airbyte_stream_with_file_transfer_support(self):
         actual_airbyte_stream = stream.as_airbyte_stream()
 
         assert actual_airbyte_stream == expected_airbyte_stream
+
+    def test_given_no_partitions_when_get_availability_then_unavailable(self) -> None:
+        self._partition_generator.generate.return_value = []
+
+        availability = self._stream.check_availability()
+
+        assert availability.is_available == False
+        assert "no stream slices were found" in availability.reason
+
+    def test_given_AirbyteTracedException_when_generating_partitions_when_get_availability_then_unavailable(self) -> None:
+        error_message = "error while generating partitions"
+        self._partition_generator.generate.side_effect = AirbyteTracedException(message=error_message)
+
+        availability = self._stream.check_availability()
+
+        assert availability.is_available == False
+        assert error_message in availability.reason
+
+    def test_given_unknown_error_when_generating_partitions_when_get_availability_then_raise(self) -> None:
+        """
+        I'm not sure why we handle AirbyteTracedException but not other exceptions but this is to keep feature compatibility with HttpAvailabilityStrategy
+        """
+        self._partition_generator.generate.side_effect = ValueError()
+        with pytest.raises(ValueError):
+            self._stream.check_availability()
+
+    def test_given_no_records_when_get_availability_then_available(self) -> None:
+        self._partition_generator.generate.return_value = [self._partition]
+        self._partition.read.return_value = []
+
+        availability = self._stream.check_availability()
+
+        assert availability.is_available == True
+
+    def test_given_records_when_get_availability_then_available(self) -> None:
+        self._partition_generator.generate.return_value = [self._partition]
+        self._partition.read.return_value = [Mock(spec=Record)]
+
+        availability = self._stream.check_availability()
+
+        assert availability.is_available == True
+
+    def test_given_AirbyteTracedException_when_reading_records_when_get_availability_then_unavailable(self) -> None:
+        self._partition_generator.generate.return_value = [self._partition]
+        error_message = "error while reading records"
+        self._partition.read.side_effect = AirbyteTracedException(message=error_message)
+
+        availability = self._stream.check_availability()
+
+        assert availability.is_available == False
+
+    def test_given_unknown_error_when_reading_record_when_get_availability_then_raise(self) -> None:
+        """
+        I'm not sure why we handle AirbyteTracedException but not other exceptions but this is to keep feature compatibility with HttpAvailabilityStrategy
+        """
+        self._partition_generator.generate.side_effect = ValueError()
+        self._partition.read.return_value = []
+        with pytest.raises(ValueError):
+            self._stream.check_availability()