"""Database engine, session factory, and startup migrations.""" from __future__ import annotations import os import sqlite3 from typing import Any from sqlalchemy import create_engine, inspect, text from sqlalchemy.orm import sessionmaker DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:////data/atc-agents.db") SQLITE_FALLBACK_PATH = os.getenv("SQLITE_FALLBACK_PATH", "/data/atc-agents.db") def make_engine(url: str = DATABASE_URL): kwargs: dict[str, Any] = {"pool_pre_ping": True} if url.startswith("sqlite"): kwargs["connect_args"] = {"check_same_thread": False} return create_engine(url, **kwargs) engine = make_engine() SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False) def migrate_approval_columns(db_engine: Any = engine) -> None: """Add approval columns on legacy SQLite/Postgres schemas.""" insp = inspect(db_engine) if "approvals" not in insp.get_table_names(): return existing = {c["name"] for c in insp.get_columns("approvals")} dialect = db_engine.dialect.name if dialect == "postgresql": alters = [ ("action_type", "VARCHAR(64)", "generic.mutate"), ("target", "TEXT", ""), ("payload", "TEXT", "{}"), ("decided_by", "VARCHAR(64)", None), ("decide_note", "TEXT", None), ("decided_at", "TIMESTAMP WITH TIME ZONE", None), ("priority", "VARCHAR(16)", "normal"), ] with db_engine.begin() as conn: for col, typ, default in alters: if col in existing: continue if default is None: conn.execute(text(f"ALTER TABLE approvals ADD COLUMN IF NOT EXISTS {col} {typ}")) elif default == "": conn.execute(text(f"ALTER TABLE approvals ADD COLUMN IF NOT EXISTS {col} {typ} DEFAULT ''")) else: conn.execute( text(f"ALTER TABLE approvals ADD COLUMN IF NOT EXISTS {col} {typ} DEFAULT '{default}'") ) return alters = [ ("action_type", "VARCHAR(64)", "'generic.mutate'"), ("target", "TEXT", "''"), ("payload", "TEXT", "'{}'"), ("decided_by", "VARCHAR(64)", "NULL"), ("decide_note", "TEXT", "NULL"), ("decided_at", "DATETIME", "NULL"), ("priority", "VARCHAR(16)", "'normal'"), ] with db_engine.begin() as conn: for col, typ, default in alters: if col not in existing: conn.execute(text(f"ALTER TABLE approvals ADD COLUMN {col} {typ} DEFAULT {default}")) def migrate_sqlite_to_postgres(db_engine: Any = engine) -> dict[str, int]: """One-time copy from legacy SQLite volume into Postgres when Postgres is empty.""" if not DATABASE_URL.startswith("postgresql"): return {} if not os.path.isfile(SQLITE_FALLBACK_PATH): return {} insp = inspect(db_engine) tables = set(insp.get_table_names()) if "approvals" not in tables or "feed" not in tables: return {} with db_engine.connect() as conn: approval_count = conn.execute(text("SELECT COUNT(*) FROM approvals")).scalar() or 0 feed_count = conn.execute(text("SELECT COUNT(*) FROM feed")).scalar() or 0 if approval_count or feed_count: return {"approvals": 0, "feed": 0, "skipped": 1} copied = {"approvals": 0, "feed": 0} src = sqlite3.connect(SQLITE_FALLBACK_PATH) src.row_factory = sqlite3.Row try: with db_engine.begin() as conn: for row in src.execute("SELECT * FROM feed"): conn.execute( text( "INSERT INTO feed (id, ts, agent_id, level, message) " "VALUES (:id, :ts, :agent_id, :level, :message) ON CONFLICT (id) DO NOTHING" ), dict(row), ) copied["feed"] += 1 for row in src.execute("SELECT * FROM approvals"): conn.execute( text( "INSERT INTO approvals (id, ts, agent_id, action, reason, status, " "action_type, target, payload, decided_by, decide_note, decided_at, priority) " "VALUES (:id, :ts, :agent_id, :action, :reason, :status, " ":action_type, :target, :payload, :decided_by, :decide_note, :decided_at, :priority) " "ON CONFLICT (id) DO NOTHING" ), { "id": row["id"], "ts": row["ts"], "agent_id": row["agent_id"], "action": row["action"], "reason": row["reason"], "status": row["status"], "action_type": row["action_type"] if "action_type" in row.keys() else "generic.mutate", "target": row["target"] if "target" in row.keys() else "", "payload": row["payload"] if "payload" in row.keys() else "{}", "decided_by": row["decided_by"] if "decided_by" in row.keys() else None, "decide_note": row["decide_note"] if "decide_note" in row.keys() else None, "decided_at": row["decided_at"] if "decided_at" in row.keys() else None, "priority": row["priority"] if "priority" in row.keys() else "normal", }, ) copied["approvals"] += 1 finally: src.close() return copied def init_database(Base: type) -> dict[str, Any]: """Create tables, run migrations, optionally import legacy SQLite data.""" Base.metadata.create_all(engine) migrate_approval_columns(engine) migrated = migrate_sqlite_to_postgres(engine) return {"engine": engine.dialect.name, "migrated_from_sqlite": migrated} def db_health() -> dict[str, Any]: try: with engine.connect() as conn: conn.execute(text("SELECT 1")) insp = inspect(engine) tables = insp.get_table_names() stats = {} if "approvals" in tables: with engine.connect() as conn: stats["approvals"] = conn.execute(text("SELECT COUNT(*) FROM approvals")).scalar() stats["approvals_pending"] = conn.execute( text("SELECT COUNT(*) FROM approvals WHERE status = 'pending'") ).scalar() if "feed" in tables: with engine.connect() as conn: stats["feed"] = conn.execute(text("SELECT COUNT(*) FROM feed")).scalar() return {"ok": True, "dialect": engine.dialect.name, "tables": tables, **stats} except Exception as exc: return {"ok": False, "dialect": engine.dialect.name, "error": str(exc)}