Update check streams to use dynamic streams configs

Author: lazebnyi

airbyte_cdk/sources/declarative/checks/__init__.py

@@ -7,7 +7,7 @@
 from pydantic.v1 import BaseModel
 
 from airbyte_cdk.sources.declarative.checks.check_dynamic_stream import CheckDynamicStream
-from airbyte_cdk.sources.declarative.checks.check_stream import CheckStream
+from airbyte_cdk.sources.declarative.checks.check_stream import CheckStream, DynamicStreamCheckConfig
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
 from airbyte_cdk.sources.declarative.models import (
     CheckDynamicStream as CheckDynamicStreamModel,
@@ -21,4 +21,4 @@
     "CheckDynamicStream": CheckDynamicStreamModel,
 }
 
-__all__ = ["CheckStream", "CheckDynamicStream", "ConnectionChecker"]
+__all__ = ["CheckStream", "CheckDynamicStream", "ConnectionChecker", "DynamicStreamCheckConfig"]

airbyte_cdk/sources/declarative/checks/check_stream.py

@@ -8,10 +8,21 @@
 from typing import Any, List, Mapping, Tuple
 
 from airbyte_cdk import AbstractSource
+from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
 from airbyte_cdk.sources.streams.http.availability_strategy import HttpAvailabilityStrategy
 
 
+@dataclass(frozen=True)
+class DynamicStreamCheckConfig:
+    """Defines the configuration for dynamic stream during connection checking. This class specifies
+    what dynamic streams  in the stream template should be updated with value, supporting dynamic interpolation
+    and type enforcement."""
+
+    dynamic_stream_name: str
+    stream_count: int = 0
+
+
 @dataclass
 class CheckStream(ConnectionChecker):
     """
@@ -22,6 +33,7 @@ class CheckStream(ConnectionChecker):
     """
 
     stream_names: List[str]
+    dynamic_streams_check_configs: List[DynamicStreamCheckConfig]
     parameters: InitVar[Mapping[str, Any]]
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
@@ -30,7 +42,15 @@ 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)
+        try:
+            streams = source.streams(config=config)
+        except Exception as error:
+            error_message = (
+                f"Encountered an error trying to connect to streams. Error: {error}"
+            )
+            logger.error(error_message, exc_info=True)
+            return False, error_message
+
         stream_name_to_stream = {s.name: s for s in streams}
         if len(streams) == 0:
             return False, f"No streams to connect to from source {source}"
@@ -53,4 +73,41 @@ def check_connection(
                     f"Encountered an error trying to connect to stream {stream_name}. Error: \n {traceback.format_exc()}"
                 )
                 return False, f"Unable to connect to stream {stream_name} - {error}"
+
+        if isinstance(source, ManifestDeclarativeSource) and self.dynamic_streams_check_configs:
+            dynamic_stream_name_to_dynamic_stream = {dynamic_stream["name"]: dynamic_stream for dynamic_stream in
+                                       source.resolved_manifest.get("dynamic_streams", [])}
+
+            dynamic_stream_name_to_generated_streams = {}
+            for stream in source.dynamic_streams:
+                dynamic_stream_name_to_generated_streams[
+                    stream["dynamic_stream_name"]] = dynamic_stream_name_to_generated_streams.setdefault(
+                    stream["dynamic_stream_name"], []).append(stream)
+
+            for dynamic_streams_check_config in self.dynamic_streams_check_configs:
+                dynamic_stream = dynamic_stream_name_to_dynamic_stream.get(dynamic_streams_check_config.dynamic_stream_name)
+
+                is_config_depend = dynamic_stream["components_resolver"]["type"] == "ConfigComponentsResolver"
+
+                if not is_config_depend and not bool(dynamic_streams_check_config.stream_count):
+                    continue
+
+                generated_streams = dynamic_stream_name_to_generated_streams.get(dynamic_streams_check_config.dynamic_stream_name)
+                availability_strategy = HttpAvailabilityStrategy()
+
+                for stream in generated_streams[: min(dynamic_streams_check_config.stream_count, len(generated_streams))]:
+                    try:
+                        stream_is_available, reason = availability_strategy.check_availability(
+                            stream, logger
+                        )
+                        if not stream_is_available:
+                            logger.warning(f"Stream {stream.name} is not available: {reason}")
+                            return False, reason
+                    except Exception as error:
+                        error_message = (
+                            f"Encountered an error trying to connect to stream {stream.name}. Error: {error}"
+                        )
+                        logger.error(error_message, exc_info=True)
+                        return False, error_message
+
         return True, None

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -316,7 +316,11 @@ definitions:
     type: object
     required:
       - type
-      - stream_names
+    anyOf:
+      - required:
+          - stream_names
+      - required:
+          - dynamic_streams_check_configs
     properties:
       type:
         type: string
@@ -330,6 +334,28 @@ definitions:
         examples:
           - ["users"]
           - ["users", "contacts"]
+      dynamic_streams_check_configs:
+        type: array
+        items:
+          "$ref": "#/definitions/DynamicStreamCheckConfig"
+  DynamicStreamCheckConfig:
+    type: object
+    required:
+      - type
+      - dynamic_stream_name
+    properties:
+      type:
+        type: string
+        enum: [ DynamicStreamCheckConfig ]
+      dynamic_stream_name:
+        title: Dynamic Stream Name
+        description: The dynamic stream name.
+        type: string
+      stream_count:
+        title: Stream Count
+        description: Numbers of the streams to try reading from when running a check operation.
+        type: integer
+        default: 0
   CheckDynamicStream:
     title: Dynamic Streams to Check
     description: (This component is experimental. Use at your own risk.) Defines the dynamic streams to try reading when running a check operation.