add check/discover

Author: pedroslopez

airbyte_cdk/manifest_runner/README.md

@@ -38,6 +38,16 @@ Test reading from a specific stream in the manifest.
 
 **POST** - Test stream reading with configurable limits for records, pages, and slices.
 
+### `/manifest/check`
+Check configuration against a manifest.
+
+**POST** - Validates connector configuration and returns success/failure status with message.
+
+### `/manifest/discover`
+Discover streams from a manifest.
+
+**POST** - Returns the catalog of available streams from the manifest.
+
 ### `/manifest/resolve` 
 Resolve a manifest to its final configuration.
 

airbyte_cdk/manifest_runner/api_models/manifest.py

@@ -7,6 +7,7 @@
 
 from typing import Any, List, Optional
 
+from airbyte_protocol_dataclasses.models import AirbyteCatalog
 from pydantic import BaseModel, Field
 
 from .dicts import ConnectorConfig, Manifest
@@ -25,6 +26,33 @@ class StreamTestReadRequest(BaseModel):
     slice_limit: int = Field(default=5, ge=1, le=20)
 
 
+class CheckRequest(BaseModel):
+    """Request to check a manifest."""
+
+    manifest: Manifest
+    config: ConnectorConfig
+
+
+class CheckResponse(BaseModel):
+    """Response to check a manifest."""
+
+    success: bool
+    message: Optional[str] = None
+
+
+class DiscoverRequest(BaseModel):
+    """Request to discover a manifest."""
+
+    manifest: Manifest
+    config: ConnectorConfig
+
+
+class DiscoverResponse(BaseModel):
+    """Response to discover a manifest."""
+
+    catalog: AirbyteCatalog
+
+
 class ResolveRequest(BaseModel):
     """Request to resolve a manifest."""
 

airbyte_cdk/manifest_runner/manifest_runner/runner.py

@@ -1,7 +1,18 @@
-from typing import Any, List, Mapping
+import logging
+from typing import Any, Iterable, List, Mapping, Tuple
+
+from airbyte_protocol_dataclasses.models import (
+    AirbyteCatalog,
+    AirbyteConnectionStatus,
+    AirbyteMessage,
+    Status,
+    TraceType,
+)
+from airbyte_protocol_dataclasses.models import Type as AirbyteMessageType
 
 from airbyte_cdk.connector_builder.models import StreamRead
 from airbyte_cdk.connector_builder.test_reader import TestReader
+from airbyte_cdk.entrypoint import AirbyteEntrypoint
 from airbyte_cdk.models.airbyte_protocol import (
     AirbyteStateMessage,
     ConfiguredAirbyteCatalog,
@@ -13,6 +24,7 @@
 
 class ManifestRunner:
     _source: ManifestDeclarativeSource
+    _logger = logging.getLogger("airbyte.manifest-runner")
 
     def __init__(self, source: ManifestDeclarativeSource) -> None:
         self._source = source
@@ -44,3 +56,85 @@ def test_read(
         )
 
         return stream_read
+
+    def check_connection(
+        self,
+        config: Mapping[str, Any],
+    ) -> Tuple[bool, str]:
+        """
+        Check the connection to the source.
+        """
+
+        spec = self._source.spec(self._logger)
+        messages = AirbyteEntrypoint(source=self._source).check(spec, config)
+        messages_by_type = self._get_messages_by_type(messages)
+        self._raise_on_trace_message(messages_by_type)
+        connection_status = self._get_connection_status(messages_by_type)
+
+        if connection_status:
+            return connection_status.status == Status.SUCCEEDED, connection_status.message
+        return False, "Connection check failed"
+
+    def discover(
+        self,
+        config: Mapping[str, Any],
+    ) -> AirbyteCatalog | None:
+        """
+        Discover the catalog from the source.
+        """
+        spec = self._source.spec(self._logger)
+        messages = AirbyteEntrypoint(source=self._source).discover(spec, config)
+        messages_by_type = self._get_messages_by_type(messages)
+        self._raise_on_trace_message(messages_by_type)
+        return self._get_catalog(messages_by_type)
+
+    def _get_messages_by_type(
+        self,
+        messages: Iterable[AirbyteMessage],
+    ) -> Mapping[str, Iterable[AirbyteMessage]]:
+        """
+        Group messages by type.
+        """
+        grouped = {}
+        for message in messages:
+            msg_type = message.type
+            if msg_type not in grouped:
+                grouped[msg_type] = []
+            grouped[msg_type].append(message)
+        return grouped
+
+    def _get_connection_status(
+        self,
+        messages_by_type: Mapping[str, List[AirbyteMessage]],
+    ) -> AirbyteConnectionStatus | None:
+        """
+        Get the connection status from the messages.
+        """
+        messages = messages_by_type.get(AirbyteMessageType.CONNECTION_STATUS)
+        return messages[-1].connectionStatus if messages else None
+
+    def _get_catalog(
+        self,
+        messages_by_type: Mapping[str, List[AirbyteMessage]],
+    ) -> AirbyteCatalog:
+        """
+        Get the catalog from the messages.
+        """
+        messages = messages_by_type.get(AirbyteMessageType.CATALOG)
+        return messages[-1].catalog if messages else None
+
+    def _raise_on_trace_message(
+        self,
+        messages_by_type: Mapping[str, List[AirbyteMessage]],
+    ) -> None:
+        """
+        Raise an exception if a trace message is found.
+        """
+        messages = [
+            message
+            for message in messages_by_type.get(AirbyteMessageType.TRACE, [])
+            if message.trace.type == TraceType.ERROR
+        ]
+        if messages:
+            # TODO: raise a better exception
+            raise Exception(messages[-1].trace.error.message)