stack/dev/scripts/generate_models.py

"""Introspect a DuckDB database and generate Pydantic table models.

Usage:
    python generate_models.py [--db PATH] [--out DIR] [--base-import IMPORT]

Defaults:
    --db  ../notebooks/aco.duckdb
    --out ../src/aco/table/
    --base-import aco.table.base
"""

from __future__ import annotations

import argparse
import re
from pathlib import Path

import duckdb

# ── DuckDB type → Python type annotation mapping ────────────────────────────

DUCKDB_TYPE_MAP: dict[str, str] = {
    "VARCHAR": "str",
    "BOOLEAN": "bool",
    "INTEGER": "int",
    "BIGINT": "int",
    "HUGEINT": "int",
    "FLOAT": "float",
    "DOUBLE": "float",
    "DATE": "date",
    "TIMESTAMP": "datetime",
}

DECIMAL_RE = re.compile(r"DECIMAL\((\d+),\s*(\d+)\)")


def python_type(duckdb_type: str) -> str:
    """Map a DuckDB column type to a Python type annotation string."""
    if duckdb_type in DUCKDB_TYPE_MAP:
        return DUCKDB_TYPE_MAP[duckdb_type]
    m = DECIMAL_RE.match(duckdb_type)
    if m:
        return "Decimal"
    return "str"  # fallback


def needs_import(py_type: str) -> set[str]:
    """Return import lines needed for a given Python type."""
    imports: set[str] = set()
    if py_type == "date":
        imports.add("from datetime import date")
    if py_type == "datetime":
        imports.add("from datetime import datetime")
    if py_type == "Decimal":
        imports.add("from decimal import Decimal")
    return imports


# ── Name helpers ─────────────────────────────────────────────────────────────


def to_class_name(schema: str, table: str) -> str:
    """Convert schema.table to PascalCase class name."""
    raw = f"{schema}__{table}"
    parts = raw.split("_")
    return "".join(p.capitalize() for p in parts if p)


def sanitize_field(name: str) -> str:
    """Ensure a column name is a valid Python identifier."""
    if name.isidentifier():
        return name
    return name.replace(" ", "_").replace("-", "_")


# ── Introspection ───────────────────────────────────────────────────────────


def get_schemas(con: duckdb.DuckDBPyConnection) -> list[str]:
    rows = con.execute(
        "SELECT DISTINCT table_schema FROM information_schema.tables "
        "ORDER BY table_schema"
    ).fetchall()
    return [r[0] for r in rows]


def get_tables(con: duckdb.DuckDBPyConnection, schema: str) -> list[str]:
    rows = con.execute(
        "SELECT table_name FROM information_schema.tables "
        "WHERE table_schema = ? ORDER BY table_name",
        [schema],
    ).fetchall()
    return [r[0] for r in rows]


def get_columns(
    con: duckdb.DuckDBPyConnection, schema: str, table: str
) -> list[tuple[str, str, bool, str]]:
    """Return [(col_name, duckdb_type, nullable, comment), ...]."""
    rows = con.execute(
        "SELECT column_name, data_type, is_nullable "
        "FROM information_schema.columns "
        "WHERE table_schema = ? AND table_name = ? "
        "ORDER BY ordinal_position",
        [schema, table],
    ).fetchall()
    cols = [(r[0], r[1], r[2] == "YES") for r in rows]

    # Try to read column comments from duckdb_columns()
    comments: dict[str, str] = {}
    try:
        crows = con.execute(
            "SELECT column_name, comment FROM duckdb_columns() "
            "WHERE schema_name = ? AND table_name = ? "
            "AND comment IS NOT NULL",
            [schema, table],
        ).fetchall()
        comments = {r[0]: r[1] for r in crows}
    except Exception:
        pass

    return [
        (name, dtype, nullable, comments.get(name, ""))
        for name, dtype, nullable in cols
    ]


# ── Code generation ─────────────────────────────────────────────────────────


def generate_model(
    schema: str, table: str, columns: list[tuple[str, str, bool, str]]
) -> str:
    """Generate a single Pydantic model class."""
    cls = to_class_name(schema, table)
    all_imports: set[str] = set()
    field_lines: list[str] = []
    uses_field = False

    for col_name, col_type, nullable, comment in columns:
        py = python_type(col_type)
        all_imports |= needs_import(py)
        fname = sanitize_field(col_name)

        if comment:
            uses_field = True
            escaped = comment.replace('"', '\\"')
            if nullable:
                field_lines.append(
                    f"    {fname}: {py} | None = Field("
                    f'default=None, description="{escaped}")'
                )
            else:
                field_lines.append(
                    f'    {fname}: {py} = Field(description="{escaped}")'
                )
        else:
            if nullable:
                field_lines.append(f"    {fname}: {py} | None = None")
            else:
                field_lines.append(f"    {fname}: {py}")

    if uses_field:
        all_imports.add("from pydantic import Field")

    lines = []
    lines.append(f"class {cls}(SQLTable):")
    lines.append(f'    """Schema: {schema} / Table: {table}"""')
    lines.append("")
    lines.append(f'    __schema__ = "{schema}"')
    lines.append(f'    __tablename__ = "{table}"')
    lines.append("")
    for fl in field_lines:
        lines.append(fl)

    return "\n".join(lines), all_imports


def generate_schema_module(
    con: duckdb.DuckDBPyConnection, schema: str, base_import: str = "aco.table.base"
) -> str:
    """Generate a full Python module for all tables in a schema."""
    tables = get_tables(con, schema)
    all_imports: set[str] = {"from __future__ import annotations"}
    all_imports.add(f"from {base_import} import SQLTable")
    models: list[str] = []

    for table in tables:
        columns = get_columns(con, schema, table)
        if not columns:
            continue
        model_code, imp = generate_model(schema, table, columns)
        all_imports |= imp
        models.append(model_code)

    # Sort imports for consistency
    sorted_imports = sorted(all_imports)
    header = "\n".join(sorted_imports)
    body = "\n\n\n".join(models)

    return f"{header}\n\n\n{body}\n"


# ── Main ────────────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate Pydantic models from DuckDB")
    from conf import cfg
    from conf import path as _conf_path

    parser.add_argument(
        "--db",
        default=str(_conf_path("db.aco")),
        help="Path to DuckDB file",
    )
    parser.add_argument(
        "--out",
        default=str(_conf_path("generate.table_out")),
        help="Output directory",
    )
    parser.add_argument(
        "--base-import",
        default=cfg.generate.base_import,
        help="Import path for the SQLTable base class",
    )
    args = parser.parse_args()

    db_path = Path(args.db)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    con = duckdb.connect(str(db_path), read_only=True)

    # Generate per-schema modules
    schemas = get_schemas(con)
    all_modules: list[str] = []

    for schema in schemas:
        safe_name = schema.replace("-", "_")
        module_path = out_dir / f"{safe_name}.py"
        code = generate_schema_module(con, schema, args.base_import)
        module_path.write_text(code)
        table_count = len(get_tables(con, schema))
        print(f"Wrote {module_path} ({table_count} tables)")
        all_modules.append(safe_name)

    # Write __init__.py
    init_lines = ['"""Generated Pydantic models for all DuckDB tables."""\n']
    for mod in sorted(all_modules):
        init_lines.append(f"from . import {mod}")
    init_lines.append("")
    (out_dir / "__init__.py").write_text("\n".join(init_lines))
    print(f"\nWrote {out_dir / '__init__.py'}")

    con.close()
    total_tables = sum(
        len(get_tables(duckdb.connect(str(db_path), read_only=True), s))
        for s in schemas
    )
    print(f"\nDone: {len(schemas)} schemas, {total_tables} tables → {out_dir}/")


if __name__ == "__main__":
    main()