pcb-afm-simulation-system/web/backend/app/database.py

"""Database connection and session management (SQLite + SQLAlchemy)."""
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker, declarative_base

from .config import DATABASE_URL

# SQLite engine (check_same_thread=False for FastAPI async)
engine = create_engine(
    DATABASE_URL,
    connect_args={"check_same_thread": False},
    echo=False,
)

# F2 fix: enable foreign key constraints for SQLite (OFF by default)
@event.listens_for(engine, "connect")
def _set_sqlite_pragma(dbapi_connection, connection_record):
    cursor = dbapi_connection.cursor()
    cursor.execute("PRAGMA foreign_keys=ON")
    cursor.close()

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


def get_db():
    """Dependency for FastAPI route database sessions."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def init_db():
    """Create all database tables and migrate existing tables."""
    from . import models  # noqa: F401 - ensure models are imported
    Base.metadata.create_all(bind=engine)
    _migrate_existing_tables()


def _migrate_existing_tables():
    """Add missing columns to existing tables (SQLite ALTER TABLE).

    SQLite doesn't support DROP COLUMN or complex migrations, but ADD COLUMN
    is supported. This handles model additions without requiring DB reset.
    """
    from sqlalchemy import inspect, text

    inspector = inspect(engine)
    if "simulation_results" not in inspector.get_table_names():
        return

    existing_cols = {c["name"] for c in inspector.get_columns("simulation_results")}
    new_cols = {
        "fidelity_level": "VARCHAR(30) DEFAULT 'L1_motorcad_emag'",
        "confidence_grade": "VARCHAR(2) DEFAULT 'C'",
        "model_template_version": "VARCHAR(100) DEFAULT ''",
        "solver_settings_hash": "VARCHAR(64) DEFAULT ''",
        "constraint_margins_json": "TEXT DEFAULT '{}'",
        "surrogate_prediction_json": "TEXT DEFAULT '{}'",
        "cross_validation_json": "TEXT DEFAULT '{}'",
        "convergence_status_json": "TEXT DEFAULT '{}'",
        "raw_json": "TEXT DEFAULT '[]'",
    }

    with engine.connect() as conn:
        for col_name, col_def in new_cols.items():
            if col_name not in existing_cols:
                conn.execute(text(f"ALTER TABLE simulation_results ADD COLUMN {col_name} {col_def}"))

    # P3-M2: migrate tasks table for adaptive-batch fields.
    if "tasks" in inspector.get_table_names():
        task_cols = {c["name"] for c in inspector.get_columns("tasks")}
        task_new = {
            "task_type": "VARCHAR(20) DEFAULT 'scan'",
            "loop_id": "VARCHAR(64)",
            "batch_id": "INTEGER",
            "point_ids": "TEXT",
            "dynamic": "INTEGER DEFAULT 0",
        }
        with engine.connect() as conn2:
            for col_name, col_def in task_new.items():
                if col_name not in task_cols:
                    conn2.execute(text(f"ALTER TABLE tasks ADD COLUMN {col_name} {col_def}"))
            conn2.commit()
        conn.commit()