atc-agents/api/db.py

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