lean into pydantic model

Author: aaronsteers

airbyte_cdk/cli/airbyte_cdk/_image.py

@@ -36,10 +36,10 @@
 
 import click
 
+from airbyte_cdk.models.connector_metadata import MetadataFile
 from airbyte_cdk.utils.docker.build import (
     build_from_base_image,
     build_from_dockerfile,
-    read_metadata,
     set_up_logging,
     verify_docker_installation,
     verify_image,
@@ -80,11 +80,11 @@ def build(
         sys.exit(1)
 
     try:
-        metadata = read_metadata(connector_directory)
-        click.echo(f"Connector: {metadata.dockerRepository}")
-        click.echo(f"Version: {metadata.dockerImageTag}")
+        metadata = MetadataFile.from_file(connector_directory / "metadata.yaml")
+        click.echo(f"Connector: {metadata.data.dockerRepository}")
+        click.echo(f"Version: {metadata.data.dockerImageTag}")
 
-        if metadata.language:
+        if metadata.data.language:
             click.echo(f"Connector language from metadata: {metadata.language}")
         else:
             click.echo("Connector language not specified in metadata")
@@ -111,7 +111,7 @@ def build(
             platforms = "linux/amd64"
             click.echo(f"Multi-platform build check failed. Building for platform: {platforms}")
 
-        if metadata.connectorBuildOptions and metadata.connectorBuildOptions.baseImage:
+        if metadata.data.connectorBuildOptions and metadata.data.connectorBuildOptions.baseImage:
             image_name = build_from_base_image(connector_dir, metadata, tag, platforms)
         else:
             image_name = build_from_dockerfile(connector_dir, metadata, tag, platforms)

airbyte_cdk/models/connector_metadata.py

@@ -3,8 +3,10 @@
 from __future__ import annotations
 
 from enum import Enum
+from pathlib import Path
 from typing import Any
 
+import yaml
 from pydantic import BaseModel, Field
 
 
@@ -54,3 +56,23 @@ class MetadataFile(BaseModel):
     model_config = {"extra": "allow"}
 
     data: ConnectorMetadata = Field(..., description="Connector metadata")
+
+    @classmethod
+    def from_file(
+        cls,
+        file_path: Path,
+    ) -> MetadataFile:
+        """Load metadata from a YAML file."""
+        if not file_path.exists():
+            raise FileNotFoundError(f"Metadata file not found: {file_path!s}")
+
+        metadata_content = file_path.read_text()
+        metadata_dict = yaml.safe_load(metadata_content)
+
+        if not metadata_dict or "data" not in metadata_dict:
+            raise ValueError(
+                "Invalid metadata format: missing 'data' field in YAML file '{file_path!s}'"
+            )
+
+        metadata_file = MetadataFile.model_validate(metadata_dict)
+        return metadata_file

airbyte_cdk/utils/docker/build.py

@@ -18,68 +18,6 @@
 logger = logging.getLogger("airbyte-cdk.utils.docker.build")
 
 
-def run_command(
-    connector_dir: Path,
-    tag: str = "dev",
-    platform: str = "linux/amd64",
-    no_verify: bool = False,
-    verbose: bool = False,
-) -> int:
-    """Run the build command with the given arguments.
-    
-    Args:
-        connector_dir: Path to the connector directory.
-        tag: Tag to apply to the built image.
-        platform: Platform to build for.
-        no_verify: Whether to skip verification of the built image.
-        verbose: Whether to enable verbose logging.
-        
-    Returns:
-        Exit code (0 for success, non-zero for failure).
-    """
-    set_up_logging(verbose)
-    
-    if not verify_docker_installation():
-        logger.error("Docker is not installed or not running. Please install Docker and try again.")
-        return 1
-    
-    if not connector_dir.exists():
-        logger.error(f"Connector directory not found: {connector_dir}")
-        return 1
-    
-    try:
-        metadata = read_metadata(connector_dir)
-        logger.info(f"Connector: {metadata.dockerRepository}")
-        logger.info(f"Version: {metadata.dockerImageTag}")
-    except (FileNotFoundError, ValueError) as e:
-        logger.error(f"Error reading connector metadata: {e}")
-        return 1
-    
-    try:
-        if metadata.connectorBuildOptions and metadata.connectorBuildOptions.baseImage:
-            image_name = build_from_base_image(connector_dir, metadata, tag, platform)
-        else:
-            image_name = build_from_dockerfile(connector_dir, metadata, tag, platform)
-        
-        if not no_verify:
-            if verify_image(image_name):
-                logger.info(f"Build completed successfully: {image_name}")
-                return 0
-            else:
-                logger.error(f"Built image failed verification: {image_name}")
-                return 1
-        else:
-            logger.info(f"Build completed successfully (without verification): {image_name}")
-            return 0
-            
-    except Exception as e:
-        logger.error(f"Error building Docker image: {e}")
-        if verbose:
-            import traceback
-            logger.error(traceback.format_exc())
-        return 1
-
-
 def set_up_logging(verbose: bool = False) -> None:
     """Set up logging configuration.
 
@@ -94,33 +32,6 @@ def set_up_logging(verbose: bool = False) -> None:
     )
 
 
-def read_metadata(connector_dir: Path) -> ConnectorMetadata:
-    """Read and parse connector metadata from metadata.yaml.
-
-    Args:
-        connector_dir: Path to the connector directory.
-
-    Returns:
-        The parsed connector metadata.
-
-    Raises:
-        FileNotFoundError: If the metadata.yaml file doesn't exist.
-        ValueError: If the metadata is invalid.
-    """
-    metadata_path = connector_dir / "metadata.yaml"
-    if not metadata_path.exists():
-        raise FileNotFoundError(f"Metadata file not found: {metadata_path}")
-
-    metadata_content = metadata_path.read_text()
-    metadata_dict = yaml.safe_load(metadata_content)
-
-    if not metadata_dict or "data" not in metadata_dict:
-        raise ValueError("Invalid metadata format: missing 'data' field")
-
-    metadata_file = MetadataFile.model_validate(metadata_dict)
-    return metadata_file.data
-
-
 def run_docker_command(cmd: List[str], check: bool = True) -> Tuple[int, str, str]:
     """Run a Docker command as a subprocess.