generate_component_manifest_files.py

#!/usr/bin/env python3
#
# Usage:
# > uv run bin/generate_component_manifest_files.py
#
# /// script
# dependencies = [
#   "datamodel-code-generator==0.26.3",
#   "PyYAML>=6.0.1",
# ]
# ///

# Copyright (c) 2024 Airbyte, Inc., all rights reserved.

import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
from glob import glob
from pathlib import Path

import yaml

LOCAL_YAML_DIR_PATH = "airbyte_cdk/sources/declarative"
LOCAL_OUTPUT_DIR_PATH = "airbyte_cdk/sources/declarative/models"


def get_all_yaml_files_without_ext() -> list[str]:
    return [Path(f).stem for f in glob(f"{LOCAL_YAML_DIR_PATH}/*.yaml")]


def generate_init_module_content() -> str:
    header = "# generated by bin/generate_component_manifest_files.py\n"
    for module_name in get_all_yaml_files_without_ext():
        header += f"from .{module_name} import *\n"
    return header


def generate_json_schema():
    """Generate JSON schema from the YAML file for schemastore.org registration.

    When registered with schemastore.org, a number of IDEs and libraries
    automatically apply the JSON Schema validation features such as:
    - auto-complete for keys and enums
    - hover-tooltips for descriptions and examples
    - linting squiggles for validation errors
    """
    yaml_file_path = f"{LOCAL_YAML_DIR_PATH}/declarative_component_schema.yaml"
    json_file_path = f"{LOCAL_YAML_DIR_PATH}/generated/declarative_component_schema.json"

    with open(yaml_file_path, "r") as yaml_file:
        schema_data = yaml.safe_load(yaml_file)

    class DateTimeEncoder(json.JSONEncoder):
        def default(self, obj):
            if hasattr(obj, "isoformat"):
                return obj.isoformat()
            return super().default(obj)

    import os

    os.makedirs(os.path.dirname(json_file_path), exist_ok=True)

    with open(json_file_path, "w") as json_file:
        json.dump(schema_data, json_file, indent=2, cls=DateTimeEncoder)

    print(f"Generated JSON schema: {json_file_path}")


def replace_base_model_for_classes_with_deprecated_fields(post_processed_content: str) -> str:
    """
    Replace the base model for classes with deprecated fields.
    This function looks for classes that inherit from `BaseModel` and have fields marked as deprecated.
    It replaces the base model with `BaseModelWithDeprecations` for those classes.
    """

    # Find classes with deprecated fields
    classes_with_deprecated_fields = set()
    class_matches = re.finditer(r"class (\w+)\(BaseModel\):", post_processed_content)

    for class_match in class_matches:
        class_name = class_match.group(1)
        class_start = class_match.start()
        # Find the next class definition or end of file
        next_class_match = re.search(
            r"class \w+\(",
            post_processed_content[class_start + len(class_match.group(0)) :],
        )
        class_end = (
            len(post_processed_content)
            if next_class_match is None
            else class_start + len(class_match.group(0)) + next_class_match.start()
        )
        class_content = post_processed_content[class_start:class_end]

        # Check if any field has deprecated=True
        if re.search(r"deprecated\s*=\s*True", class_content):
            classes_with_deprecated_fields.add(class_name)

    # update the imports to include the new base model with deprecation warinings
    # only if there are classes with the fields marked as deprecated.
    if len(classes_with_deprecated_fields) > 0:
        # Find where to insert the base model - after imports but before class definitions
        imports_end = post_processed_content.find(
            "\n\n",
            post_processed_content.find("from pydantic.v1 import"),
        )
        if imports_end > 0:
            post_processed_content = (
                post_processed_content[:imports_end]
                + "\n\n"
                + "from airbyte_cdk.sources.declarative.models.base_model_with_deprecations import (\n"
                + "    BaseModelWithDeprecations,\n"
                + ")"
                + post_processed_content[imports_end:]
            )

    # Use the `BaseModelWithDeprecations` base model for the classes with deprecated fields
    for class_name in classes_with_deprecated_fields:
        pattern = rf"class {class_name}\(BaseModel\):"
        replacement = f"class {class_name}(BaseModelWithDeprecations):"
        post_processed_content = re.sub(pattern, replacement, post_processed_content)

    return post_processed_content


def post_process_codegen(generated_dir: str, post_processed_dir: str):
    """Post-process generated files to fix pydantic imports and deprecated fields."""
    os.makedirs(post_processed_dir, exist_ok=True)

    for generated_file in os.listdir(generated_dir):
        if generated_file.endswith(".py"):
            input_path = os.path.join(generated_dir, generated_file)
            output_path = os.path.join(post_processed_dir, generated_file)

            with open(input_path, "r") as f:
                original_content = f.read()

            # the space before _parameters is intentional to avoid replacing things like `request_parameters:` with `requestparameters:`
            post_processed_content = original_content.replace(
                " _parameters:", " parameters:"
            ).replace("from pydantic", "from pydantic.v1")

            post_processed_content = replace_base_model_for_classes_with_deprecated_fields(
                post_processed_content
            )

            with open(output_path, "w") as f:
                f.write(post_processed_content)


def main():
    generate_json_schema()
    init_module_content = generate_init_module_content()

    with tempfile.TemporaryDirectory() as temp_dir:
        generated_dir = os.path.join(temp_dir, "generated")
        post_processed_dir = os.path.join(temp_dir, "generated_post_processed")

        os.makedirs(generated_dir, exist_ok=True)

        init_file_path = os.path.join(generated_dir, "__init__.py")
        with open(init_file_path, "w") as f:
            f.write(init_module_content)

        for yaml_file in get_all_yaml_files_without_ext():
            input_yaml = os.path.join(LOCAL_YAML_DIR_PATH, f"{yaml_file}.yaml")
            output_py = os.path.join(generated_dir, f"{yaml_file}.py")

            cmd = [
                "datamodel-codegen",
                "--input",
                input_yaml,
                "--output",
                output_py,
                "--disable-timestamp",
                "--enum-field-as-literal",
                "one",
                "--set-default-enum-member",
                "--use-double-quotes",
                "--remove-special-field-name-prefix",
                # allow usage of the extra key such as `deprecated`, etc.
                "--field-extra-keys",
                # account the `deprecated` flag provided for the field.
                "deprecated",
                # account the `deprecation_message` provided for the field.
                "deprecation_message",
            ]

            try:
                result = subprocess.run(cmd, check=True, capture_output=True, text=True)
                print(f"Generated {output_py}")
            except subprocess.CalledProcessError as e:
                print(f"Error generating {output_py}: {e}")
                print(f"stdout: {e.stdout}")
                print(f"stderr: {e.stderr}")
                sys.exit(1)

        post_process_codegen(generated_dir, post_processed_dir)

        if os.path.exists(LOCAL_OUTPUT_DIR_PATH):
            shutil.rmtree(LOCAL_OUTPUT_DIR_PATH)
        shutil.copytree(post_processed_dir, LOCAL_OUTPUT_DIR_PATH)

        print(f"Generated models exported to {LOCAL_OUTPUT_DIR_PATH}")


if __name__ == "__main__":
    main()