scripts/generate_models_from_schema.py

"""
Script to download Spider schema and auto-generate SQLAlchemy models.

The spider-schema dataset contains detailed database schemas including
table names, column names, types, and relationships. This script
downloads the schema and generates SQLAlchemy ORM models.

Usage:
    # Generate models for student_assessment database
    python generate_models_from_schema.py --db-id student_assessment

    # Generate for multiple databases
    python generate_models_from_schema.py --db-id all --output-dir models/

    # Load from validation split
    python generate_models_from_schema.py --db-id student_assessment --split validation
"""

import argparse
from pathlib import Path
from typing import Any, Dict, List
from datasets import load_dataset


# Type mapping from Spider schema to SQLAlchemy
SQLALCHEMY_TYPE_MAP = {
    "number": "Integer",
    "int": "Integer",
    "float": "Float",
    "text": "String",
    "string": "String",
    "varchar": "String",
    "char": "String",
    "date": "Date",
    "datetime": "DateTime",
    "timestamp": "DateTime",
    "time": "DateTime",
    "boolean": "Boolean",
    "bool": "Boolean",
}


def get_sqlalchemy_type(col_type: str) -> str:
    """Convert Spider schema type to SQLAlchemy type."""
    col_type_lower = col_type.lower().strip()

    # Exact match
    if col_type_lower in SQLALCHEMY_TYPE_MAP:
        return SQLALCHEMY_TYPE_MAP[col_type_lower]

    # Substring match (e.g., "varchar(255)" -> "String")
    for key, sa_type in SQLALCHEMY_TYPE_MAP.items():
        if key in col_type_lower:
            return sa_type

    # Default to String
    return "String"


def generate_model_code(
    db_id: str,
    tables: List[Dict[str, Any]],
    schema: Dict[str, Any],
) -> str:
    """Generate SQLAlchemy model code from schema.

    Args:
        db_id: Database ID
        tables: List of table schemas
        schema: Full schema dictionary with relationships

    Returns:
        Generated Python code as string
    """
    sa_import = (
        "from sqlalchemy import Column, Integer, String,"
        " Float, Date, DateTime, Boolean, ForeignKey"
    )
    lines = [
        '"""',
        f"SQLAlchemy ORM models for '{db_id}' database.",
        "",
        "Auto-generated from Spider schema dataset.",
        '"""',
        "",
        "from datetime import datetime",
        sa_import,
        "from sqlalchemy.ext.declarative import declarative_base",
        "from sqlalchemy.orm import relationship",
        "",
        "Base = declarative_base()",
        "",
    ]

    # Generate model for each table
    table_names = [t["name"] for t in tables]

    for table in tables:
        table_name = table["name"]
        class_name = "".join(word.capitalize() for word in table_name.split("_"))

        lines.append(f"class {class_name}(Base):")
        lines.append(f'    """Model for {table_name} table."""')
        lines.append(f'    __tablename__ = "{table_name}"')
        lines.append("")

        # Add columns
        columns = table.get("columns", [])
        for col in columns:
            col_name = col["name"]
            col_type = col.get("type", "text")
            sa_type = get_sqlalchemy_type(col_type)

            # Determine if primary key
            is_pk = col.get("is_primary_key", False)

            # Determine if foreign key
            fk_str = ""
            for fk in schema.get("foreign_keys", []):
                tbl_idx = table_names.index(table_name)
                col_idx = columns.index(col)
                if fk[0] == (tbl_idx, col_idx):
                    _source_idx, target_table_idx = fk
                    target_col_idx = fk[2] if len(fk) > 2 else 0
                    target_table = table_names[target_table_idx]
                    target_col = tables[target_table_idx]["columns"][target_col_idx][
                        "name"
                    ]
                    fk_str = f', ForeignKey("{target_table}.{target_col}")'

            # Default nullable to False for primary keys
            nullable = "False" if is_pk else "True"
            pk_str = ", primary_key=True" if is_pk else ""

            type_arg = ""
            if sa_type == "String" and "(" in col_type:
                type_arg = f"({col_type.split('(')[1].rstrip(')')})"

            lines.append(
                f"    {col_name} = Column("
                f"{sa_type}{type_arg}"
                f"{pk_str}{fk_str}, nullable={nullable})"
            )

        lines.append("")

    return "\n".join(lines)


def download_schema_and_generate_models(
    db_id: str = "student_assessment",
    split: str = "train",
    output_dir: str = "data/models",
) -> None:
    """Download Spider schema and generate SQLAlchemy models.

    Args:
        db_id: Database ID to download schema for
        split: Dataset split ("train" or "validation")
        output_dir: Directory to save generated model files
    """
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    print(f"Loading Spider schema dataset ({split} split)...")
    dataset = load_dataset("richardr1126/spider-schema", split=split)

    if db_id.lower() == "all":
        # Generate models for all databases
        processed = set()
        for item in dataset:
            current_db_id = item.get("db_id")
            if current_db_id in processed:
                continue
            processed.add(current_db_id)

            tables = item.get("table", [])

            # Generate code (simplified)
            code = generate_simplified_models(current_db_id, tables)

            filepath = output_path / f"{current_db_id}.py"
            with open(filepath, "w") as f:
                f.write(code)

            print(f"  {current_db_id}: {len(tables)} tables → {filepath}")
    else:
        # Filter for specific db_id
        matching = [item for item in dataset if item.get("db_id") == db_id]

        if not matching:
            print(f"No schema found for db_id='{db_id}'")
            return

        item = matching[0]
        tables = item.get("table", [])

        # Generate simplified model code
        code = generate_simplified_models(db_id, tables)

        filepath = output_path / f"{db_id}.py"
        with open(filepath, "w") as f:
            f.write(code)

        print(f"Found schema for db_id='{db_id}' with {len(tables)} tables")
        print(f"Generated models → {filepath}")
        print(f"\nTables: {', '.join(t['name'] for t in tables)}")


def generate_simplified_models(db_id: str, tables: List[Dict[str, Any]]) -> str:
    """Generate SQLAlchemy models from table schema (simplified version).

    Args:
        db_id: Database ID
        tables: List of table definitions from schema

    Returns:
        Generated Python code
    """
    sa_import = (
        "from sqlalchemy import Column, Integer, String,"
        " Float, Date, DateTime, Boolean, ForeignKey"
    )
    lines = [
        '"""',
        f"SQLAlchemy ORM models for '{db_id}' database.",
        "",
        "Auto-generated from Spider schema dataset.",
        '"""',
        "",
        "from datetime import datetime",
        sa_import,
        "from sqlalchemy.ext.declarative import declarative_base",
        "from sqlalchemy.orm import relationship",
        "",
        "Base = declarative_base()",
        "",
    ]

    for table in tables:
        table_name = table.get("name", "Unknown")
        class_name = "".join(word.capitalize() for word in table_name.split("_"))

        lines.append("")
        lines.append(f"class {class_name}(Base):")
        lines.append(f'    """Model for {table_name} table."""')
        lines.append(f'    __tablename__ = "{table_name}"')
        lines.append("")

        # Add columns
        columns = table.get("columns", [])
        if columns:
            for col in columns:
                col_name = col.get("name", "unknown")
                col_type = col.get("type", "text")
                sa_type = get_sqlalchemy_type(col_type)

                # Determine string length from type
                length_spec = ""
                if sa_type == "String":
                    if "(" in col_type and ")" in col_type:
                        length = col_type.split("(")[1].split(")")[0]
                        if length.isdigit():
                            length_spec = f"({length})"
                    else:
                        length_spec = "(255)"  # default

                lines.append(
                    f"    {col_name} = Column({sa_type}{length_spec}, nullable=True)"
                )
        else:
            lines.append("    id = Column(Integer, primary_key=True)")

        lines.append("")

    return "\n".join(lines)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Download Spider schema and generate SQLAlchemy models",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--db-id",
        type=str,
        default="student_assessment",
        help="Database ID to generate models for (or 'all' for all databases)",
    )
    parser.add_argument(
        "--split",
        type=str,
        default="train",
        choices=["train", "validation"],
        help="Schema dataset split to use",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="data/models",
        help="Directory to save generated model files",
    )

    args = parser.parse_args()
    download_schema_and_generate_models(
        db_id=args.db_id, split=args.split, output_dir=args.output_dir
    )