"""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()