concurrent cdk

Author: pedroslopez

airbyte_cdk/manifest_server/app.py

@@ -3,7 +3,7 @@
 from ..manifest_server.routers import capabilities, health, manifest
 
 app = FastAPI(
-    title="Manifest Server Service",
+    title="Manifest Server",
     description="A service for running low-code Airbyte connectors",
     version="0.1.0",
     contact={

airbyte_cdk/manifest_server/command_processor/utils.py

@@ -1,4 +1,6 @@
-from typing import Any, Mapping, Optional
+from typing import Any, List, Mapping, Optional
+
+from airbyte_protocol_dataclasses.models import AirbyteStateMessage
 
 from airbyte_cdk.models import (
     AirbyteStream,
@@ -7,6 +9,10 @@
     DestinationSyncMode,
     SyncMode,
 )
+from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
+    ConcurrentDeclarativeSource,
+    TestLimits,
+)
 from airbyte_cdk.sources.declarative.manifest_declarative_source import (
     ManifestDeclarativeSource,
 )
@@ -56,21 +62,31 @@ def should_normalize_manifest(manifest: Mapping[str, Any]) -> bool:
 
 def build_source(
     manifest: Mapping[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,
-) -> ManifestDeclarativeSource:
-    return ManifestDeclarativeSource(
+) -> 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.
+    if "concurrency_level" in manifest:
+        manifest["concurrency_level"]["default_concurrency"] = 1
+    else:
+        manifest["concurrency_level"] = {"type": "ConcurrencyLevel", "default_concurrency": 1}
+
+    return ConcurrentDeclarativeSource(
+        catalog=catalog,
+        state=state,
         source_config=manifest,
         config=config,
         normalize_manifest=should_normalize_manifest(manifest),
         migrate_manifest=should_migrate_manifest(manifest),
         emit_connector_builder_messages=True,
-        component_factory=ModelToComponentFactory(
-            emit_connector_builder_messages=True,
-            limit_pages_fetched_per_slice=page_limit,
-            limit_slices_fetched=slice_limit,
-            disable_retries=True,
-            disable_cache=True,
+        limits=TestLimits(
+            max_pages_per_slice=page_limit,
+            max_slices=slice_limit,
+            max_records=record_limit,
         ),
     )

airbyte_cdk/manifest_server/routers/manifest.py

@@ -3,6 +3,7 @@
 from typing import Any, Dict, List, Mapping, Optional
 
 import jsonschema
+from airbyte_protocol_dataclasses.models import AirbyteStateMessage, ConfiguredAirbyteCatalog
 from fastapi import APIRouter, Depends, HTTPException
 
 from airbyte_cdk.manifest_server.api_models.manifest import (
@@ -12,7 +13,9 @@
     DiscoverResponse,
 )
 from airbyte_cdk.models import AirbyteStateMessageSerializer
-from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
+from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
+    ConcurrentDeclarativeSource,
+)
 from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
     INJECTED_COMPONENTS_PY,
     INJECTED_COMPONENTS_PY_CHECKSUMS,
@@ -34,12 +37,23 @@
 def safe_build_source(
     manifest_dict: Mapping[str, Any],
     config_dict: Mapping[str, Any],
+    catalog: Optional[ConfiguredAirbyteCatalog] = None,
+    state: Optional[List[AirbyteStateMessage]] = None,
     page_limit: Optional[int] = None,
     slice_limit: Optional[int] = None,
-) -> ManifestDeclarativeSource:
+    record_limit: Optional[int] = None,
+) -> ConcurrentDeclarativeSource[Optional[List[AirbyteStateMessage]]]:
     """Wrapper around build_source that converts ValidationError to HTTPException."""
     try:
-        return build_source(manifest_dict, config_dict, page_limit, slice_limit)
+        return build_source(
+            manifest_dict,
+            catalog,
+            config_dict,
+            state,
+            record_limit,
+            page_limit,
+            slice_limit,
+        )
     except jsonschema.exceptions.ValidationError as e:
         raise HTTPException(status_code=400, detail=f"Invalid manifest: {e.message}")
 
@@ -58,10 +72,16 @@ def test_read(request: StreamTestReadRequest) -> StreamRead:
     """
     config_dict = request.config.model_dump()
 
+    catalog = build_catalog(request.stream_name)
     source = safe_build_source(
-        request.manifest.model_dump(), config_dict, request.page_limit, request.slice_limit
+        request.manifest.model_dump(),
+        config_dict,
+        catalog,
+        request.state,
+        request.page_limit,
+        request.slice_limit,
+        request.record_limit,
     )
-    catalog = build_catalog(request.stream_name)
     state = [AirbyteStateMessageSerializer.load(state) for state in request.state]
 
     if request.custom_components_code:

unit_tests/manifest_server/command_processor/test_utils.py

@@ -31,15 +31,12 @@ def test_build_catalog_creates_correct_structure(self):
         assert configured_stream.sync_mode == SyncMode.incremental
         assert configured_stream.destination_sync_mode == DestinationSyncMode.overwrite
 
-    @patch("airbyte_cdk.manifest_server.command_processor.utils.ManifestDeclarativeSource")
-    @patch("airbyte_cdk.manifest_server.command_processor.utils.ModelToComponentFactory")
+    @patch("airbyte_cdk.manifest_server.command_processor.utils.ConcurrentDeclarativeSource")
     def test_build_source_creates_manifest_declarative_source(
-        self, mock_component_factory_class, mock_source_class
+        self, mock_source_class
     ):
-        """Test that build_source creates a ManifestDeclarativeSource with correct parameters."""
+        """Test that build_source creates a ConcurrentDeclarativeSource with correct parameters."""
         # Setup mocks
-        mock_component_factory = Mock()
-        mock_component_factory_class.return_value = mock_component_factory
         mock_source = Mock()
         mock_source_class.return_value = mock_source
 
@@ -68,64 +65,57 @@ def test_build_source_creates_manifest_declarative_source(
             "timeout": 30,
         }
 
-        # Call build_source
-        result = build_source(manifest, config)
+        # Call build_source with additional parameters
+        catalog = build_catalog("test_stream")
+        state = []
+        result = build_source(manifest, catalog, config, state)
 
-        # Verify ModelToComponentFactory was created with correct parameters
-        mock_component_factory_class.assert_called_once_with(
-            emit_connector_builder_messages=True,
-            limit_pages_fetched_per_slice=None,
-            limit_slices_fetched=None,
-            disable_retries=True,
-            disable_cache=True,
-        )
-
-        # Verify ManifestDeclarativeSource was created with correct parameters
+        # Verify ConcurrentDeclarativeSource was created with correct parameters
         mock_source_class.assert_called_once_with(
+            catalog=catalog,
+            state=state,
             source_config=manifest,
             config=config,
             normalize_manifest=False,  # Default when flag not set
             migrate_manifest=False,  # Default when flag not set
             emit_connector_builder_messages=True,
-            component_factory=mock_component_factory,
+            limits=mock_source_class.call_args[1]["limits"],
         )
 
         assert result == mock_source
 
-    @patch("airbyte_cdk.manifest_server.command_processor.utils.ManifestDeclarativeSource")
-    @patch("airbyte_cdk.manifest_server.command_processor.utils.ModelToComponentFactory")
+    @patch("airbyte_cdk.manifest_server.command_processor.utils.ConcurrentDeclarativeSource")
     def test_build_source_with_normalize_flag(
-        self, mock_component_factory_class, mock_source_class
+        self, mock_source_class
     ):
         """Test build_source when normalize flag is set."""
-        mock_component_factory = Mock()
-        mock_component_factory_class.return_value = mock_component_factory
         mock_source = Mock()
         mock_source_class.return_value = mock_source
 
         manifest = {"streams": [{"name": "test_stream"}], SHOULD_NORMALIZE_KEY: True}
         config = {"api_key": "test_key"}
+        catalog = build_catalog("test_stream")
+        state = []
 
-        build_source(manifest, config)
+        build_source(manifest, catalog, config, state)
 
         # Verify normalize_manifest is True
         call_args = mock_source_class.call_args[1]
         assert call_args["normalize_manifest"] is True
         assert call_args["migrate_manifest"] is False
 
-    @patch("airbyte_cdk.manifest_server.command_processor.utils.ManifestDeclarativeSource")
-    @patch("airbyte_cdk.manifest_server.command_processor.utils.ModelToComponentFactory")
-    def test_build_source_with_migrate_flag(self, mock_component_factory_class, mock_source_class):
+    @patch("airbyte_cdk.manifest_server.command_processor.utils.ConcurrentDeclarativeSource")
+    def test_build_source_with_migrate_flag(self, mock_source_class):
         """Test build_source when migrate flag is set."""
-        mock_component_factory = Mock()
-        mock_component_factory_class.return_value = mock_component_factory
         mock_source = Mock()
         mock_source_class.return_value = mock_source
 
         manifest = {"streams": [{"name": "test_stream"}], SHOULD_MIGRATE_KEY: True}
         config = {"api_key": "test_key"}
+        catalog = build_catalog("test_stream")
+        state = []
 
-        build_source(manifest, config)
+        build_source(manifest, catalog, config, state)
 
         # Verify migrate_manifest is True
         call_args = mock_source_class.call_args[1]