Merge branch 'main' into lazebnyi/update-config-components-resolver-to-support-list-of-stream-configs

Author: lazebnyi

.coveragerc

@@ -9,3 +9,13 @@ omit =
 
     # omit as unimplemented
     airbyte_cdk/base_python/cdk/streams/auth/jwt.py
+
+    # omit temporary files and test files
+    /tmp/tmp*.py
+    unit_tests/*
+
+[paths]
+# Reconcile file paths
+source =
+    ./
+    /tmp/

airbyte_cdk/cli/source_declarative_manifest/_run.py

@@ -16,15 +16,17 @@
 
 from __future__ import annotations
 
+import argparse
 import json
 import pkgutil
 import sys
 import traceback
-from collections.abc import Mapping
+from collections.abc import MutableMapping
 from pathlib import Path
 from typing import Any, cast
 
 import orjson
+import yaml
 
 from airbyte_cdk.entrypoint import AirbyteEntrypoint, launch
 from airbyte_cdk.models import (
@@ -54,8 +56,9 @@ class SourceLocalYaml(YamlDeclarativeSource):
     def __init__(
         self,
         catalog: ConfiguredAirbyteCatalog | None,
-        config: Mapping[str, Any] | None,
+        config: MutableMapping[str, Any] | None,
         state: TState,
+        config_path: str | None = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -74,6 +77,7 @@ def __init__(
             config=config,
             state=state,  # type: ignore [arg-type]
             path_to_yaml="manifest.yaml",
+            config_path=config_path,
         )
 
 
@@ -91,8 +95,14 @@ def handle_command(args: list[str]) -> None:
 
 def _get_local_yaml_source(args: list[str]) -> SourceLocalYaml:
     try:
-        config, catalog, state = _parse_inputs_into_config_catalog_state(args)
-        return SourceLocalYaml(config=config, catalog=catalog, state=state)
+        parsed_args = AirbyteEntrypoint.parse_args(args)
+        config, catalog, state = _parse_inputs_into_config_catalog_state(parsed_args)
+        return SourceLocalYaml(
+            config=config,
+            catalog=catalog,
+            state=state,
+            config_path=parsed_args.config if hasattr(parsed_args, "config") else None,
+        )
     except Exception as error:
         print(
             orjson.dumps(
@@ -162,26 +172,46 @@ def create_declarative_source(
     connector builder.
     """
     try:
-        config: Mapping[str, Any] | None
+        config: MutableMapping[str, Any] | None
         catalog: ConfiguredAirbyteCatalog | None
         state: list[AirbyteStateMessage]
-        config, catalog, state = _parse_inputs_into_config_catalog_state(args)
-        if config is None or "__injected_declarative_manifest" not in config:
+
+        parsed_args = AirbyteEntrypoint.parse_args(args)
+        config, catalog, state = _parse_inputs_into_config_catalog_state(parsed_args)
+
+        if config is None:
             raise ValueError(
                 "Invalid config: `__injected_declarative_manifest` should be provided at the root "
-                f"of the config but config only has keys: {list(config.keys() if config else [])}"
+                "of the config or using the --manifest-path argument."
+            )
+
+        # If a manifest_path is provided in the args, inject it into the config
+        if hasattr(parsed_args, "manifest_path") and parsed_args.manifest_path:
+            injected_manifest = _parse_manifest_from_file(parsed_args.manifest_path)
+            if injected_manifest:
+                config["__injected_declarative_manifest"] = injected_manifest
+
+        if "__injected_declarative_manifest" not in config:
+            raise ValueError(
+                "Invalid config: `__injected_declarative_manifest` should be provided at the root "
+                "of the config or using the --manifest-path argument. "
+                f"Config only has keys: {list(config.keys() if config else [])}"
             )
         if not isinstance(config["__injected_declarative_manifest"], dict):
             raise ValueError(
                 "Invalid config: `__injected_declarative_manifest` should be a dictionary, "
                 f"but got type: {type(config['__injected_declarative_manifest'])}"
             )
 
+        if hasattr(parsed_args, "components_path") and parsed_args.components_path:
+            _register_components_from_file(parsed_args.components_path)
+
         return ConcurrentDeclarativeSource(
             config=config,
             catalog=catalog,
             state=state,
             source_config=cast(dict[str, Any], config["__injected_declarative_manifest"]),
+            config_path=parsed_args.config if hasattr(parsed_args, "config") else None,
         )
     except Exception as error:
         print(
@@ -205,13 +235,12 @@ def create_declarative_source(
 
 
 def _parse_inputs_into_config_catalog_state(
-    args: list[str],
+    parsed_args: argparse.Namespace,
 ) -> tuple[
-    Mapping[str, Any] | None,
+    MutableMapping[str, Any] | None,
     ConfiguredAirbyteCatalog | None,
     list[AirbyteStateMessage],
 ]:
-    parsed_args = AirbyteEntrypoint.parse_args(args)
     config = (
         ConcurrentDeclarativeSource.read_config(parsed_args.config)
         if hasattr(parsed_args, "config")
@@ -231,6 +260,44 @@ def _parse_inputs_into_config_catalog_state(
     return config, catalog, state
 
 
+def _parse_manifest_from_file(filepath: str) -> dict[str, Any] | None:
+    """Extract and parse a manifest file specified in the args."""
+    try:
+        with open(filepath, "r", encoding="utf-8") as manifest_file:
+            manifest_content = yaml.safe_load(manifest_file)
+            if manifest_content is None:
+                raise ValueError(f"Manifest file at {filepath} is empty")
+            if not isinstance(manifest_content, dict):
+                raise ValueError(f"Manifest must be a dictionary, got {type(manifest_content)}")
+            return manifest_content
+    except Exception as error:
+        raise ValueError(f"Failed to load manifest file from {filepath}: {error}")
+
+
+def _register_components_from_file(filepath: str) -> None:
+    """Load and register components from a Python file specified in the args."""
+    import importlib.util
+    import sys
+
+    components_path = Path(filepath)
+
+    module_name = "components"
+    sdm_module_name = "source_declarative_manifest.components"
+
+    # Create module spec
+    spec = importlib.util.spec_from_file_location(module_name, components_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Could not load module from {components_path}")
+
+    # Create module and execute code, registering the module before executing its code
+    # To avoid issues with dataclasses that look up the module
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    sys.modules[sdm_module_name] = module
+
+    spec.loader.exec_module(module)
+
+
 def run() -> None:
     args: list[str] = sys.argv[1:]
     handle_command(args)

airbyte_cdk/connector.py

@@ -8,7 +8,7 @@
 import os
 import pkgutil
 from abc import ABC, abstractmethod
-from typing import Any, Generic, Mapping, Optional, Protocol, TypeVar
+from typing import Any, Generic, Mapping, MutableMapping, Optional, Protocol, TypeVar
 
 import yaml
 
@@ -41,9 +41,9 @@ def configure(self, config: Mapping[str, Any], temp_dir: str) -> TConfig:
         """
 
     @staticmethod
-    def read_config(config_path: str) -> Mapping[str, Any]:
+    def read_config(config_path: str) -> MutableMapping[str, Any]:
         config = BaseConnector._read_json_file(config_path)
-        if isinstance(config, Mapping):
+        if isinstance(config, MutableMapping):
             return config
         else:
             raise ValueError(

airbyte_cdk/entrypoint.py

@@ -84,6 +84,18 @@ def parse_args(args: List[str]) -> argparse.Namespace:
         required_check_parser.add_argument(
             "--config", type=str, required=True, help="path to the json configuration file"
         )
+        check_parser.add_argument(
+            "--manifest-path",
+            type=str,
+            required=False,
+            help="path to the YAML manifest file to inject into the config",
+        )
+        check_parser.add_argument(
+            "--components-path",
+            type=str,
+            required=False,
+            help="path to the custom components file, if it exists",
+        )
 
         # discover
         discover_parser = subparsers.add_parser(
@@ -95,6 +107,18 @@ def parse_args(args: List[str]) -> argparse.Namespace:
         required_discover_parser.add_argument(
             "--config", type=str, required=True, help="path to the json configuration file"
         )
+        discover_parser.add_argument(
+            "--manifest-path",
+            type=str,
+            required=False,
+            help="path to the YAML manifest file to inject into the config",
+        )
+        discover_parser.add_argument(
+            "--components-path",
+            type=str,
+            required=False,
+            help="path to the custom components file, if it exists",
+        )
 
         # read
         read_parser = subparsers.add_parser(
@@ -114,6 +138,18 @@ def parse_args(args: List[str]) -> argparse.Namespace:
             required=True,
             help="path to the catalog used to determine which data to read",
         )
+        read_parser.add_argument(
+            "--manifest-path",
+            type=str,
+            required=False,
+            help="path to the YAML manifest file to inject into the config",
+        )
+        read_parser.add_argument(
+            "--components-path",
+            type=str,
+            required=False,
+            help="path to the custom components file, if it exists",
+        )
 
         return main_parser.parse_args(args)
 

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -74,6 +74,7 @@ def __init__(
         debug: bool = False,
         emit_connector_builder_messages: bool = False,
         component_factory: Optional[ModelToComponentFactory] = None,
+        config_path: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
         # todo: We could remove state from initialization. Now that streams are grouped during the read(), a source
@@ -96,6 +97,7 @@ def __init__(
             debug=debug,
             emit_connector_builder_messages=emit_connector_builder_messages,
             component_factory=component_factory,
+            config_path=config_path,
         )
 
         concurrency_level_from_manifest = self._source_config.get("concurrency_level")