move to EntrypointOutput

Author: pedroslopez

airbyte_cdk/manifest_server/command_processor/processor.py

@@ -1,14 +1,10 @@
 import logging
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple
+from typing import Any, List, Mapping, Optional, Tuple
 
 from airbyte_protocol_dataclasses.models import (
     AirbyteCatalog,
-    AirbyteConnectionStatus,
-    AirbyteMessage,
     Status,
-    TraceType,
 )
-from airbyte_protocol_dataclasses.models import Type as AirbyteMessageType
 from fastapi import HTTPException
 
 from airbyte_cdk.connector_builder.models import StreamRead
@@ -18,9 +14,11 @@
     AirbyteStateMessage,
     ConfiguredAirbyteCatalog,
 )
+from airbyte_cdk.models.airbyte_protocol_serializers import AirbyteMessageSerializer
 from airbyte_cdk.sources.declarative.manifest_declarative_source import (
     ManifestDeclarativeSource,
 )
+from airbyte_cdk.test.entrypoint_wrapper import AirbyteEntrypointException, EntrypointOutput
 
 
 class ManifestCommandProcessor:
@@ -63,20 +61,29 @@ def test_read(
     def check_connection(
         self,
         config: Mapping[str, Any],
-    ) -> Tuple[bool, str]:
+    ) -> Tuple[bool, Optional[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)
+        entrypoint = AirbyteEntrypoint(source=self._source)
+        messages = entrypoint.check(spec, config)
+        output = EntrypointOutput(
+            messages=[AirbyteEntrypoint.airbyte_message_to_string(m) for m in messages],
+            command=["check"],
+        )
+        self._raise_on_trace_message(output)
+
+        status_messages = output.connection_status_messages
+        if not status_messages or status_messages[-1].connectionStatus is None:
+            return False, "Connection check did not return a status message"
 
-        if connection_status:
-            return connection_status.status == Status.SUCCEEDED, connection_status.message
-        return False, "Connection check failed"
+        connection_status = status_messages[-1].connectionStatus
+        return (
+            connection_status.status == Status.SUCCEEDED,
+            connection_status.message,
+        )
 
     def discover(
         self,
@@ -86,66 +93,29 @@ def discover(
         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],
-    ) -> Dict[str, List[AirbyteMessage]]:
-        """
-        Group messages by type.
-        """
-        grouped: Dict[str, List[AirbyteMessage]] = {}
-        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]],
-    ) -> Optional[AirbyteConnectionStatus]:
-        """
-        Get the connection status from the messages.
-        """
-        messages = messages_by_type.get(AirbyteMessageType.CONNECTION_STATUS)
-        return messages[-1].connectionStatus if messages else None
+        entrypoint = AirbyteEntrypoint(source=self._source)
+        messages = entrypoint.discover(spec, config)
+        output = EntrypointOutput(
+            messages=[AirbyteEntrypoint.airbyte_message_to_string(m) for m in messages],
+            command=["discover"],
+        )
+        self._raise_on_trace_message(output)
 
-    def _get_catalog(
-        self,
-        messages_by_type: Mapping[str, List[AirbyteMessage]],
-    ) -> Optional[AirbyteCatalog]:
-        """
-        Get the catalog from the messages.
-        """
-        messages = messages_by_type.get(AirbyteMessageType.CATALOG)
-        return messages[-1].catalog if messages else None
+        try:
+            catalog_message = output.catalog
+            return catalog_message.catalog
+        except ValueError:
+            # No catalog message found
+            return None
 
     def _raise_on_trace_message(
         self,
-        messages_by_type: Mapping[str, List[AirbyteMessage]],
+        output: EntrypointOutput,
     ) -> 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:
-            error_message = messages[-1].trace.error.message
-            self._logger.warning(
-                "Error response from CDK: %s\n%s",
-                error_message,
-                messages[-1].trace.error.stack_trace,
-            )
-            raise HTTPException(
-                status_code=422,
-                detail=f"AirbyteTraceMessage response from CDK: {error_message}",
-            )
+        try:
+            output.raise_if_errors()
+        except AirbyteEntrypointException as e:
+            raise HTTPException(status_code=422, detail=e.message)