airbyte-cdk 6.61.2__py3-none-any.whl → 6.61.3.post2.dev17299502224__py3-none-any.whl

airbyte_cdk/manifest_server/command_processor/processor.py

@@ -0,0 +1,122 @@
+import logging
+from typing import Any, List, Mapping, Optional, Tuple
+
+from airbyte_protocol_dataclasses.models import (
+    AirbyteCatalog,
+    Status,
+)
+from fastapi import HTTPException
+
+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 import (
+    AirbyteStateMessage,
+    ConfiguredAirbyteCatalog,
+)
+from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
+    ConcurrentDeclarativeSource,
+)
+from airbyte_cdk.test.entrypoint_wrapper import AirbyteEntrypointException, EntrypointOutput
+
+
+class ManifestCommandProcessor:
+    _source: ConcurrentDeclarativeSource[Optional[List[AirbyteStateMessage]]]
+    _logger = logging.getLogger("airbyte.manifest-server")
+
+    def __init__(
+        self, source: ConcurrentDeclarativeSource[Optional[List[AirbyteStateMessage]]]
+    ) -> None:
+        self._source = source
+
+    def test_read(
+        self,
+        config: Mapping[str, Any],
+        catalog: ConfiguredAirbyteCatalog,
+        state: List[AirbyteStateMessage],
+        record_limit: int,
+        page_limit: int,
+        slice_limit: int,
+    ) -> StreamRead:
+        """
+        Test the read method of the source.
+        """
+
+        test_read_handler = TestReader(
+            max_pages_per_slice=page_limit,
+            max_slices=slice_limit,
+            max_record_limit=record_limit,
+        )
+
+        stream_read = test_read_handler.run_test_read(
+            source=self._source,
+            config=config,
+            configured_catalog=catalog,
+            state=state,
+            stream_name=catalog.streams[0].stream.name,
+            record_limit=record_limit,
+        )
+
+        return stream_read
+
+    def check_connection(
+        self,
+        config: Mapping[str, Any],
+    ) -> Tuple[bool, Optional[str]]:
+        """
+        Check the connection to the source.
+        """
+
+        spec = self._source.spec(self._logger)
+        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"
+
+        connection_status = status_messages[-1].connectionStatus
+        return (
+            connection_status.status == Status.SUCCEEDED,
+            connection_status.message,
+        )
+
+    def discover(
+        self,
+        config: Mapping[str, Any],
+    ) -> Optional[AirbyteCatalog]:
+        """
+        Discover the catalog from the source.
+        """
+        spec = self._source.spec(self._logger)
+        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)
+
+        try:
+            catalog_message = output.catalog
+            return catalog_message.catalog
+        except ValueError:
+            # No catalog message found
+            return None
+
+    def _raise_on_trace_message(
+        self,
+        output: EntrypointOutput,
+    ) -> None:
+        """
+        Raise an exception if a trace message is found.
+        """
+        try:
+            output.raise_if_errors()
+        except AirbyteEntrypointException as e:
+            raise HTTPException(status_code=422, detail=e.message)

airbyte_cdk/manifest_server/command_processor/utils.py

@@ -0,0 +1,99 @@
+import copy
+from typing import Any, Dict, List, Mapping, Optional
+
+from airbyte_protocol_dataclasses.models import AirbyteStateMessage
+
+from airbyte_cdk.models import (
+    AirbyteStream,
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteStream,
+    DestinationSyncMode,
+    SyncMode,
+)
+from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
+    ConcurrentDeclarativeSource,
+    TestLimits,
+)
+
+SHOULD_NORMALIZE_KEY = "__should_normalize"
+SHOULD_MIGRATE_KEY = "__should_migrate"
+
+
+def build_catalog(stream_name: str) -> ConfiguredAirbyteCatalog:
+    return ConfiguredAirbyteCatalog(
+        streams=[
+            ConfiguredAirbyteStream(
+                stream=AirbyteStream(
+                    name=stream_name,
+                    json_schema={},
+                    supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental],
+                ),
+                sync_mode=SyncMode.incremental,
+                destination_sync_mode=DestinationSyncMode.overwrite,
+            )
+        ]
+    )
+
+
+def should_migrate_manifest(manifest: Mapping[str, Any]) -> bool:
+    """
+    Determines whether the manifest should be migrated,
+    based on the presence of the "__should_migrate" key.
+
+    This flag is set by the UI.
+    """
+    return manifest.get(SHOULD_MIGRATE_KEY, False)
+
+
+def should_normalize_manifest(manifest: Mapping[str, Any]) -> bool:
+    """
+    Determines whether the manifest should be normalized,
+    based on the presence of the "__should_normalize" key.
+
+    This flag is set by the UI.
+    """
+    return manifest.get(SHOULD_NORMALIZE_KEY, False)
+
+
+def build_source(
+    manifest: Dict[str, Any],
+    catalog: Optional[ConfiguredAirbyteCatalog],
+    config: Mapping[str, Any],
+    state: Optional[List[AirbyteStateMessage]],
+    record_limit: Optional[int] = None,
+    page_limit: Optional[int] = None,
+    slice_limit: Optional[int] = None,
+) -> ConcurrentDeclarativeSource[Optional[List[AirbyteStateMessage]]]:
+    # We enforce a concurrency level of 1 so that the stream is processed on a single thread
+    # to retain ordering for the grouping of the builder message responses.
+    definition = copy.deepcopy(manifest)
+    if "concurrency_level" in definition:
+        definition["concurrency_level"]["default_concurrency"] = 1
+    else:
+        definition["concurrency_level"] = {
+            "type": "ConcurrencyLevel",
+            "default_concurrency": 1,
+        }
+
+    should_normalize = should_normalize_manifest(manifest)
+    if should_normalize:
+        del definition[SHOULD_NORMALIZE_KEY]
+
+    should_migrate = should_migrate_manifest(manifest)
+    if should_migrate:
+        del definition[SHOULD_MIGRATE_KEY]
+
+    return ConcurrentDeclarativeSource(
+        catalog=catalog,
+        state=state,
+        source_config=definition,
+        config=config,
+        normalize_manifest=should_normalize,
+        migrate_manifest=should_migrate,
+        emit_connector_builder_messages=True,
+        limits=TestLimits(
+            max_pages_per_slice=page_limit or TestLimits.DEFAULT_MAX_PAGES_PER_SLICE,
+            max_slices=slice_limit or TestLimits.DEFAULT_MAX_SLICES,
+            max_records=record_limit or TestLimits.DEFAULT_MAX_RECORDS,
+        ),
+    )

airbyte_cdk/manifest_server/main.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Main entry point for the Airbyte Manifest Server server."""
+
+import uvicorn
+
+
+def run_server(
+    host: str = "127.0.0.1", port: int = 8000, reload: bool = False, log_level: str = "info"
+) -> None:
+    """Run the FastAPI server."""
+
+    print(f"🚀 Starting Airbyte CDK Manifest Server on {host}:{port}")
+
+    uvicorn.run(
+        "airbyte_cdk.manifest_server.app:app",
+        host=host,
+        port=port,
+        reload=reload,
+        log_level=log_level,
+    )
+
+
+if __name__ == "__main__":
+    run_server()