cleveragents-core/scripts/create_template_db.py

#!/usr/bin/env python3
"""Create a pre-migrated template SQLite database for fast test setup.

Instead of running 25 Alembic migrations per scenario (~0.5-3s each),
tests can copy this template file (~1ms) and get an identical schema.

Usage:
    python scripts/create_template_db.py [output_path]

Default output: build/.template-migrated.db
"""

import os
import sys
from pathlib import Path

# Ensure src/ is importable
src_dir = Path(__file__).resolve().parent.parent / "src"
sys.path.insert(0, str(src_dir))


def create_template(output_path: str = "build/.template-migrated.db") -> None:
    """Create a fully-migrated template SQLite database.

    Uses Base.metadata.create_all() to create all tables in a single DDL
    batch (~5ms), then stamps the alembic_version table with the head
    revision so MigrationRunner sees no pending migrations.
    """
    from alembic import command
    from alembic.config import Config
    from alembic.script import ScriptDirectory
    from sqlalchemy import create_engine

    from cleveragents.infrastructure.database.models import Base

    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)

    # Remove existing template so we always create fresh
    if out.exists():
        out.unlink()

    db_url = f"sqlite:///{out.resolve()}"
    engine = create_engine(db_url, connect_args={"check_same_thread": False})

    # Create all 33 tables in one fast DDL batch
    Base.metadata.create_all(engine)

    # Stamp alembic_version with the head revision so MigrationRunner
    # sees the database as fully migrated (no pending migrations).
    alembic_ini = Path(__file__).resolve().parent.parent / "alembic.ini"
    cfg = Config(str(alembic_ini))
    sd = ScriptDirectory.from_config(cfg)
    head = sd.get_current_head()
    if head is None:
        msg = "No Alembic revisions found — cannot stamp template database."
        raise RuntimeError(msg)

    with engine.connect() as conn:
        cfg.attributes["connection"] = conn
        command.stamp(cfg, head)
        conn.commit()

    engine.dispose()

    # Ensure the template database is writable (0o664 = rw-rw-r--)
    # This prevents sqlite3.OperationalError: attempt to write a readonly database
    # when tests copy and modify the template during test setup.
    os.chmod(str(out), 0o664)


if __name__ == "__main__":
    path = sys.argv[1] if len(sys.argv) > 1 else "build/.template-migrated.db"
    create_template(path)