Merge branch 'main' into brian/property_chunking

Author: brianjlai

airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -4,7 +4,7 @@
 
 
 from dataclasses import asdict, dataclass, field
-from typing import Any, List, Mapping
+from typing import Any, Dict, List, Mapping
 
 from airbyte_cdk.connector_builder.test_reader import TestReader
 from airbyte_cdk.models import (
@@ -27,30 +27,34 @@
 DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE = 5
 DEFAULT_MAXIMUM_NUMBER_OF_SLICES = 5
 DEFAULT_MAXIMUM_RECORDS = 100
+DEFAULT_MAXIMUM_STREAMS = 100
 
 MAX_PAGES_PER_SLICE_KEY = "max_pages_per_slice"
 MAX_SLICES_KEY = "max_slices"
 MAX_RECORDS_KEY = "max_records"
+MAX_STREAMS_KEY = "max_streams"
 
 
 @dataclass
-class TestReadLimits:
+class TestLimits:
     max_records: int = field(default=DEFAULT_MAXIMUM_RECORDS)
     max_pages_per_slice: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE)
     max_slices: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_SLICES)
+    max_streams: int = field(default=DEFAULT_MAXIMUM_STREAMS)
 
 
-def get_limits(config: Mapping[str, Any]) -> TestReadLimits:
+def get_limits(config: Mapping[str, Any]) -> TestLimits:
     command_config = config.get("__test_read_config", {})
     max_pages_per_slice = (
         command_config.get(MAX_PAGES_PER_SLICE_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE
     )
     max_slices = command_config.get(MAX_SLICES_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_SLICES
     max_records = command_config.get(MAX_RECORDS_KEY) or DEFAULT_MAXIMUM_RECORDS
-    return TestReadLimits(max_records, max_pages_per_slice, max_slices)
+    max_streams = command_config.get(MAX_STREAMS_KEY) or DEFAULT_MAXIMUM_STREAMS
+    return TestLimits(max_records, max_pages_per_slice, max_slices, max_streams)
 
 
-def create_source(config: Mapping[str, Any], limits: TestReadLimits) -> ManifestDeclarativeSource:
+def create_source(config: Mapping[str, Any], limits: TestLimits) -> ManifestDeclarativeSource:
     manifest = config["__injected_declarative_manifest"]
     return ManifestDeclarativeSource(
         config=config,
@@ -71,7 +75,7 @@ def read_stream(
     config: Mapping[str, Any],
     configured_catalog: ConfiguredAirbyteCatalog,
     state: List[AirbyteStateMessage],
-    limits: TestReadLimits,
+    limits: TestLimits,
 ) -> AirbyteMessage:
     try:
         test_read_handler = TestReader(
@@ -117,13 +121,23 @@ def resolve_manifest(source: ManifestDeclarativeSource) -> AirbyteMessage:
         return error.as_airbyte_message()
 
 
-def full_resolve_manifest(source: ManifestDeclarativeSource) -> AirbyteMessage:
+def full_resolve_manifest(source: ManifestDeclarativeSource, limits: TestLimits) -> AirbyteMessage:
     try:
         manifest = {**source.resolved_manifest}
         streams = manifest.get("streams", [])
         for stream in streams:
             stream["dynamic_stream_name"] = None
-        streams.extend(source.dynamic_streams)
+
+        mapped_streams: Dict[str, List[Dict[str, Any]]] = {}
+        for stream in source.dynamic_streams:
+            generated_streams = mapped_streams.setdefault(stream["dynamic_stream_name"], [])
+
+            if len(generated_streams) < limits.max_streams:
+                generated_streams += [stream]
+
+        for generated_streams_list in mapped_streams.values():
+            streams.extend(generated_streams_list)
+
         manifest["streams"] = streams
         return AirbyteMessage(
             type=Type.RECORD,

airbyte_cdk/connector_builder/main.py

@@ -10,7 +10,7 @@
 
 from airbyte_cdk.connector import BaseConnector
 from airbyte_cdk.connector_builder.connector_builder_handler import (
-    TestReadLimits,
+    TestLimits,
     create_source,
     full_resolve_manifest,
     get_limits,
@@ -73,7 +73,7 @@ def handle_connector_builder_request(
     config: Mapping[str, Any],
     catalog: Optional[ConfiguredAirbyteCatalog],
     state: List[AirbyteStateMessage],
-    limits: TestReadLimits,
+    limits: TestLimits,
 ) -> AirbyteMessage:
     if command == "resolve_manifest":
         return resolve_manifest(source)
@@ -83,7 +83,7 @@ def handle_connector_builder_request(
         ), "`test_read` requires a valid `ConfiguredAirbyteCatalog`, got None."
         return read_stream(source, config, catalog, state, limits)
     elif command == "full_resolve_manifest":
-        return full_resolve_manifest(source)
+        return full_resolve_manifest(source, limits)
     else:
         raise ValueError(f"Unrecognized command {command}.")
 

airbyte_cdk/sources/declarative/checks/__init__.py

@@ -7,7 +7,10 @@
 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 +24,4 @@
     "CheckDynamicStream": CheckDynamicStreamModel,
 }
 
-__all__ = ["CheckStream", "CheckDynamicStream", "ConnectionChecker"]
+__all__ = ["CheckStream", "CheckDynamicStream", "ConnectionChecker", "DynamicStreamCheckConfig"]

airbyte_cdk/sources/declarative/checks/check_stream.py

@@ -5,13 +5,23 @@
 import logging
 import traceback
 from dataclasses import InitVar, dataclass
-from typing import Any, List, Mapping, Tuple
+from typing import Any, Dict, List, Mapping, Optional, Tuple
 
 from airbyte_cdk import AbstractSource
 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):
     """
@@ -23,34 +33,126 @@ class CheckStream(ConnectionChecker):
 
     stream_names: List[str]
     parameters: InitVar[Mapping[str, Any]]
+    dynamic_streams_check_configs: Optional[List[DynamicStreamCheckConfig]] = None
 
     def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         self._parameters = parameters
+        if self.dynamic_streams_check_configs is None:
+            self.dynamic_streams_check_configs = []
+
+    def _log_error(self, logger: logging.Logger, action: str, error: Exception) -> Tuple[bool, str]:
+        """Logs an error and returns a formatted error message."""
+        error_message = f"Encountered an error while {action}. Error: {error}"
+        logger.error(error_message + f"Error traceback: \n {traceback.format_exc()}", exc_info=True)
+        return False, error_message
 
     def check_connection(
         self, source: AbstractSource, logger: logging.Logger, config: Mapping[str, Any]
     ) -> Tuple[bool, Any]:
-        streams = source.streams(config=config)
+        """Checks the connection to the source and its streams."""
+        try:
+            streams = source.streams(config=config)
+            if not streams:
+                return False, f"No streams to connect to from source {source}"
+        except Exception as error:
+            return self._log_error(logger, "discovering streams", error)
+
         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}"
         for stream_name in self.stream_names:
-            if stream_name not in stream_name_to_stream.keys():
+            if stream_name not in stream_name_to_stream:
                 raise ValueError(
-                    f"{stream_name} is not part of the catalog. Expected one of {stream_name_to_stream.keys()}."
+                    f"{stream_name} is not part of the catalog. Expected one of {list(stream_name_to_stream.keys())}."
                 )
 
+            stream_availability, message = self._check_stream_availability(
+                stream_name_to_stream, stream_name, logger
+            )
+            if not stream_availability:
+                return stream_availability, message
+
+        should_check_dynamic_streams = (
+            hasattr(source, "resolved_manifest")
+            and hasattr(source, "dynamic_streams")
+            and self.dynamic_streams_check_configs
+        )
+
+        if should_check_dynamic_streams:
+            return self._check_dynamic_streams_availability(source, stream_name_to_stream, logger)
+
+        return True, None
+
+    def _check_stream_availability(
+        self, stream_name_to_stream: Dict[str, Any], 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]
-            availability_strategy = HttpAvailabilityStrategy()
+            stream_is_available, reason = availability_strategy.check_availability(stream, logger)
+            if not stream_is_available:
+                message = f"Stream {stream_name} is not available: {reason}"
+                logger.warning(message)
+                return stream_is_available, message
+        except Exception as error:
+            return self._log_error(logger, f"checking availability of stream {stream_name}", error)
+        return True, None
+
+    def _check_dynamic_streams_availability(
+        self, source: AbstractSource, stream_name_to_stream: Dict[str, Any], 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
+        dynamic_stream_name_to_dynamic_stream = {
+            ds.get("name", f"dynamic_stream_{i}"): ds for i, ds in enumerate(dynamic_streams)
+        }
+        generated_streams = self._map_generated_streams(source.dynamic_streams)  # type: ignore[attr-defined] # The source's dynamic_streams manifest is checked before calling this method
+
+        for check_config in self.dynamic_streams_check_configs:  # type: ignore[union-attr] # None value for self.dynamic_streams_check_configs handled in __post_init__
+            if check_config.dynamic_stream_name not in dynamic_stream_name_to_dynamic_stream:
+                return (
+                    False,
+                    f"Dynamic stream {check_config.dynamic_stream_name} is not found in manifest.",
+                )
+
+            generated = generated_streams.get(check_config.dynamic_stream_name, [])
+            stream_availability, message = self._check_generated_streams_availability(
+                generated, stream_name_to_stream, logger, check_config.stream_count
+            )
+            if not stream_availability:
+                return stream_availability, message
+
+        return True, None
+
+    def _map_generated_streams(
+        self, dynamic_streams: List[Dict[str, Any]]
+    ) -> Dict[str, List[Dict[str, Any]]]:
+        """Maps dynamic stream names to their corresponding generated streams."""
+        mapped_streams: Dict[str, List[Dict[str, Any]]] = {}
+        for stream in dynamic_streams:
+            mapped_streams.setdefault(stream["dynamic_stream_name"], []).append(stream)
+        return mapped_streams
+
+    def _check_generated_streams_availability(
+        self,
+        generated_streams: List[Dict[str, Any]],
+        stream_name_to_stream: Dict[str, Any],
+        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
                 )
                 if not stream_is_available:
-                    return False, reason
+                    message = f"Dynamic Stream {stream.name} is not available: {reason}"
+                    logger.warning(message)
+                    return False, message
             except Exception as error:
-                logger.error(
-                    f"Encountered an error trying to connect to stream {stream_name}. Error: \n {traceback.format_exc()}"
+                return self._log_error(
+                    logger, f"checking availability of dynamic stream {stream.name}", error
                 )
-                return False, f"Unable to connect to stream {stream_name} - {error}"
         return True, None

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -316,7 +316,6 @@ definitions:
     type: object
     required:
       - type
-      - stream_names
     properties:
       type:
         type: string
@@ -330,6 +329,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: The number of streams to attempt reading from during a check operation. If `stream_count` exceeds the total number of available streams, the minimum of the two values will be used.
+        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.

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -397,6 +397,12 @@ def _dynamic_stream_configs(
             for dynamic_stream in components_resolver.resolve_components(
                 stream_template_config=stream_template_config
             ):
+                dynamic_stream = {
+                    **ManifestComponentTransformer().propagate_types_and_parameters(
+                        "", dynamic_stream, {}, use_parent_parameters=True
+                    )
+                }
+
                 if "type" not in dynamic_stream:
                     dynamic_stream["type"] = "DeclarativeStream"
 

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -44,13 +44,15 @@ class BearerAuthenticator(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
-class CheckStream(BaseModel):
-    type: Literal["CheckStream"]
-    stream_names: List[str] = Field(
-        ...,
-        description="Names of the streams to try reading from when running a check operation.",
-        examples=[["users"], ["users", "contacts"]],
-        title="Stream Names",
+class DynamicStreamCheckConfig(BaseModel):
+    type: Literal["DynamicStreamCheckConfig"]
+    dynamic_stream_name: str = Field(
+        ..., description="The dynamic stream name.", title="Dynamic Stream Name"
+    )
+    stream_count: Optional[int] = Field(
+        0,
+        description="Numbers of the streams to try reading from when running a check operation.",
+        title="Stream Count",
     )
 
 
@@ -1568,6 +1570,17 @@ class AuthFlow(BaseModel):
     oauth_config_specification: Optional[OAuthConfigSpecification] = None
 
 
+class CheckStream(BaseModel):
+    type: Literal["CheckStream"]
+    stream_names: Optional[List[str]] = Field(
+        None,
+        description="Names of the streams to try reading from when running a check operation.",
+        examples=[["users"], ["users", "contacts"]],
+        title="Stream Names",
+    )
+    dynamic_streams_check_configs: Optional[List[DynamicStreamCheckConfig]] = None
+
+
 class IncrementingCountCursor(BaseModel):
     type: Literal["IncrementingCountCursor"]
     cursor_field: str = Field(