Add Command Center v2: DQ/RAG integration, S3 browser, Jupyter, GPU matrix.
Mirror mo/atc-GPU layout with config/, docs/, scripts/ for Gitea deploy.
This commit is contained in:
mo
@@ -0,0 +1,164 @@
|
||||
"""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)}
|
||||
Reference in New Issue
Block a user