atc-agents/api/sql_console.py

"""Live SQL console — PostgreSQL, MySQL, MongoDB + Trino with benchmark."""

from __future__ import annotations

import os
import re
import time
from typing import Any

import httpx
import psycopg2
import pymysql
from fastapi import APIRouter
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from pymongo import MongoClient

DB_HOST = os.getenv("DB_VAULT_HOST", "10.0.21.51")
PG_PORT = int(os.getenv("PG_PORT", "5432"))
PG_USER = os.getenv("PG_USER", "mo")
PG_PASS = os.getenv("PG_PASSWORD", "Dell2026!")
PG_DB = os.getenv("PG_DATABASE", "postgres")

MYSQL_PORT = int(os.getenv("MYSQL_PORT", "3306"))
MYSQL_USER = os.getenv("MYSQL_USER", "mo")
MYSQL_PASS = os.getenv("MYSQL_PASSWORD", "Dell2026!")
MYSQL_DB = os.getenv("MYSQL_DATABASE", "hr")

MONGO_HOST = os.getenv("MONGO_HOST", DB_HOST)
MONGO_PORT = int(os.getenv("MONGO_PORT", "27017"))
MONGO_DB = os.getenv("MONGO_DATABASE", "supplychain")

TRINO_URL = os.getenv("TRINO_URL", "http://10.0.21.50:8089").rstrip("/")
TRINO_USER = os.getenv("TRINO_USER", "atc")

router = APIRouter(prefix="/api/sql", tags=["sql"])

SAMPLES: dict[str, list[dict[str, str]]] = {
    "postgres": [
        {"id": "pg1", "label": "Server version", "sql": "SELECT version();"},
        {"id": "pg2", "label": "Current session", "sql": "SELECT current_database(), current_user, now();"},
        {"id": "pg3", "label": "Public tables", "sql": "SELECT tablename FROM pg_tables WHERE schemaname='public' ORDER BY 1 LIMIT 20;"},
        {"id": "pg4", "label": "Table count", "sql": "SELECT count(*) AS public_tables FROM information_schema.tables WHERE table_schema='public';"},
        {"id": "pg5", "label": "Active queries", "sql": "SELECT pid, usename, state, left(query, 100) AS query FROM pg_stat_activity WHERE state='active' AND pid <> pg_backend_pid() LIMIT 10;"},
        {"id": "pg6", "label": "Database sizes", "sql": "SELECT datname, pg_size_pretty(pg_database_size(datname)) AS size FROM pg_database ORDER BY pg_database_size(datname) DESC LIMIT 10;"},
        {"id": "pg7", "label": "Row counts (stats)", "sql": "SELECT relname, n_live_tup FROM pg_stat_user_tables ORDER BY n_live_tup DESC NULLS LAST LIMIT 10;"},
        {"id": "pg8", "label": "Connections", "sql": "SELECT count(*) AS connections, state FROM pg_stat_activity GROUP BY state;"},
        {"id": "pg9", "label": "Explain scan 100k", "sql": "EXPLAIN ANALYZE SELECT count(*) FROM generate_series(1, 100000);"},
        {"id": "pg10", "label": "Memory settings", "sql": "SELECT name, setting, unit FROM pg_settings WHERE name IN ('max_connections','shared_buffers','work_mem','effective_cache_size');"},
    ],
    "mysql": [
        {"id": "my1", "label": "Server version", "sql": "SELECT VERSION() AS version;"},
        {"id": "my2", "label": "Current session", "sql": "SELECT DATABASE() AS db, USER() AS user, NOW() AS now_ts;"},
        {"id": "my3", "label": "HR tables", "sql": "SHOW TABLES;"},
        {"id": "my4", "label": "Table sizes", "sql": "SELECT table_name, table_rows, ROUND((data_length+index_length)/1024/1024,1) AS mb FROM information_schema.tables WHERE table_schema='hr' ORDER BY (data_length+index_length) DESC;"},
        {"id": "my5", "label": "Employee events sample", "sql": "SELECT event_type, COUNT(*) AS cnt FROM employee_events GROUP BY event_type ORDER BY cnt DESC LIMIT 10;"},
        {"id": "my6", "label": "Recent events", "sql": "SELECT employee_id, event_type, event_date FROM employee_events ORDER BY event_date DESC LIMIT 10;"},
        {"id": "my7", "label": "Active connections", "sql": "SHOW STATUS LIKE 'Threads_connected';"},
        {"id": "my8", "label": "InnoDB buffer pool", "sql": "SHOW STATUS LIKE 'Innodb_buffer_pool%';"},
        {"id": "my9", "label": "Explain count scan", "sql": "EXPLAIN SELECT COUNT(*) FROM employee_events;"},
        {"id": "my10", "label": "Events per employee", "sql": "SELECT employee_id, COUNT(*) AS events FROM employee_events GROUP BY employee_id ORDER BY events DESC LIMIT 10;"},
    ],
    "mongodb": [
        {"id": "mg1", "label": "List databases", "sql": "SHOW DATABASES"},
        {"id": "mg2", "label": "Supplychain collections", "sql": "SHOW COLLECTIONS supplychain"},
        {"id": "mg3", "label": "Count events", "sql": "COUNT supplychain.events"},
        {"id": "mg4", "label": "Sample document", "sql": "FIND supplychain.events LIMIT 3"},
        {"id": "mg5", "label": "Events by type", "sql": "AGGREGATE supplychain.events GROUP type TOP 10"},
        {"id": "mg6", "label": "Events by region", "sql": "AGGREGATE supplychain.events GROUP region TOP 10"},
        {"id": "mg7", "label": "Distinct event types", "sql": "DISTINCT supplychain.events type"},
        {"id": "mg8", "label": "Server status", "sql": "SERVER STATUS"},
        {"id": "mg9", "label": "Indexes on events", "sql": "INDEXES supplychain.events"},
        {"id": "mg10", "label": "Recent events", "sql": "FIND supplychain.events SORT timestamp DESC LIMIT 5"},
    ],
    "trino": [
        {"id": "tq1", "label": "Trino version", "sql": "SELECT version()"},
        {"id": "tq2", "label": "Catalogs", "sql": "SHOW CATALOGS"},
        {"id": "tq3", "label": "Iceberg schemas", "sql": "SHOW SCHEMAS FROM iceberg"},
        {"id": "tq4", "label": "Sales tables (federated PG)", "sql": "SHOW TABLES FROM postgres_sales.public"},
        {"id": "tq5", "label": "Federated PG catalog", "sql": "SHOW SCHEMAS FROM postgres_sales"},
        {"id": "tq6", "label": "Cluster nodes", "sql": "SELECT node_id, http_uri, state FROM system.runtime.nodes"},
        {"id": "tq7", "label": "Running queries", "sql": "SELECT query_id, state, user FROM system.runtime.queries WHERE state != 'FINISHED' LIMIT 10"},
        {"id": "tq8", "label": "Count 1M rows (parallel)", "sql": "SELECT count(*) AS rows FROM unnest(sequence(1,1000)) a(x) CROSS JOIN unnest(sequence(1,1000)) b(y)"},
        {"id": "tq9", "label": "Federated sales by region", "sql": "SELECT region, count(*) AS orders FROM postgres_sales.public.sales_orders WHERE order_id < 500000 GROUP BY region ORDER BY orders DESC"},
        {"id": "tq10", "label": "Parallel sin() 1M rows", "sql": "SELECT sum(sin(a.x)) AS total FROM unnest(sequence(1,1000)) a(x) CROSS JOIN unnest(sequence(1,1000)) b(y)"},
    ],
}

BENCHMARK_SQL = {
    "postgres": "SELECT sum(sin(i)) AS total FROM generate_series(1, 5000000) i",
    "trino": "SELECT sum(sin(x + (y - 1) * 5000)) AS total FROM unnest(sequence(1,5000)) a(x) CROSS JOIN unnest(sequence(1,1000)) b(y)",
}


def _tabular(columns: list[str], rows: list[list[Any]], elapsed_ms: int, **extra: Any) -> dict[str, Any]:
    return {
        "ok": True,
        "columns": columns,
        "rows": rows,
        "row_count": len(rows),
        "elapsed_ms": elapsed_ms,
        "truncated": extra.pop("truncated", False),
        **extra,
    }


def _run_postgres(sql: str, limit: int = 200) -> dict[str, Any]:
    t0 = time.perf_counter()
    conn = psycopg2.connect(
        host=DB_HOST, port=PG_PORT, user=PG_USER, password=PG_PASS, dbname=PG_DB, connect_timeout=8,
    )
    try:
        conn.set_session(readonly=True, autocommit=True)
        cur = conn.cursor()
        cur.execute(sql)
        elapsed_ms = int((time.perf_counter() - t0) * 1000)
        if cur.description:
            columns = [d[0] for d in cur.description]
            rows = [list(r) for r in cur.fetchmany(limit)]
            return _tabular(columns, rows, elapsed_ms, truncated=len(rows) >= limit)
        return _tabular([], [], elapsed_ms, message="OK")
    finally:
        conn.close()


def _run_mysql(sql: str, limit: int = 200) -> dict[str, Any]:
    t0 = time.perf_counter()
    conn = pymysql.connect(
        host=DB_HOST, port=MYSQL_PORT, user=MYSQL_USER, password=MYSQL_PASS,
        database=MYSQL_DB, connect_timeout=8, read_timeout=30,
    )
    try:
        cur = conn.cursor()
        cur.execute(sql)
        elapsed_ms = int((time.perf_counter() - t0) * 1000)
        if cur.description:
            columns = [d[0] for d in cur.description]
            rows = [list(r) for r in cur.fetchmany(limit)]
            return _tabular(columns, rows, elapsed_ms, truncated=len(rows) >= limit)
        return _tabular([], [], elapsed_ms, message="OK")
    finally:
        conn.close()


def _mongo_client() -> MongoClient:
    return MongoClient(f"mongodb://{MONGO_HOST}:{MONGO_PORT}/", serverSelectionTimeoutMS=8000)


def _fmt_mongo_value(val: Any) -> Any:
    if val is None:
        return None
    if isinstance(val, (str, int, float, bool)):
        return val
    return str(val)[:200]


def _run_mongo(query: str, limit: int = 200) -> dict[str, Any]:
    t0 = time.perf_counter()
    q = query.strip()
    client = _mongo_client()
    try:
        columns: list[str] = []
        rows: list[list[Any]] = []

        if q == "SHOW DATABASES":
            columns = ["database"]
            rows = [[name] for name in client.list_database_names()]
        elif m := re.match(r"^SHOW COLLECTIONS(?:\s+(\w+))?$", q, re.I):
            db_name = m.group(1) or MONGO_DB
            columns = ["collection"]
            rows = [[name] for name in client[db_name].list_collection_names()]
        elif m := re.match(r"^COUNT\s+(\w+)\.(\w+)$", q, re.I):
            db_name, coll = m.group(1), m.group(2)
            columns = ["count"]
            rows = [[client[db_name][coll].estimated_document_count()]]
        elif m := re.match(r"^FIND\s+(\w+)\.(\w+)(?:\s+SORT\s+(\w+)\s+DESC)?(?:\s+LIMIT\s+(\d+))?$", q, re.I):
            db_name, coll, sort_field, lim = m.group(1), m.group(2), m.group(3), m.group(4)
            n = min(int(lim) if lim else 10, limit)
            cursor = client[db_name][coll].find({})
            if sort_field:
                cursor = cursor.sort(sort_field, -1)
            docs = list(cursor.limit(n))
            if not docs:
                columns = ["result"]
                rows = [["(empty)"]]
            else:
                columns = sorted({k for d in docs for k in d if k != "_id"})
                rows = [[_fmt_mongo_value(d.get(c)) for c in columns] for d in docs]
        elif m := re.match(r"^DISTINCT\s+(\w+)\.(\w+)\s+(\w+)$", q, re.I):
            db_name, coll, field = m.group(1), m.group(2), m.group(3)
            columns = [field]
            vals = client[db_name][coll].distinct(field)[:limit]
            rows = [[_fmt_mongo_value(v)] for v in vals]
        elif m := re.match(r"^AGGREGATE\s+(\w+)\.(\w+)\s+GROUP\s+(\w+)\s+TOP\s+(\d+)$", q, re.I):
            db_name, coll, field, top_n = m.group(1), m.group(2), m.group(3), int(m.group(4))
            pipe = [
                {"$group": {"_id": f"${field}", "count": {"$sum": 1}}},
                {"$sort": {"count": -1}},
                {"$limit": min(top_n, limit)},
            ]
            columns = [field, "count"]
            rows = [[r.get("_id"), r.get("count")] for r in client[db_name][coll].aggregate(pipe, maxTimeMS=15000)]
        elif m := re.match(r"^INDEXES\s+(\w+)\.(\w+)$", q, re.I):
            db_name, coll = m.group(1), m.group(2)
            columns = ["name", "keys"]
            rows = [[i.get("name"), str(i.get("key"))] for i in client[db_name][coll].list_indexes()]
        elif q.upper() == "SERVER STATUS":
            status = client.admin.command("serverStatus")
            columns = ["metric", "value"]
            rows = [
                ["version", status.get("version")],
                ["uptime_secs", status.get("uptime")],
                ["connections_current", status.get("connections", {}).get("current")],
                ["mem_resident_mb", status.get("mem", {}).get("resident")],
            ]
        else:
            return {"ok": False, "error": f"Unknown MongoDB command: {q[:120]}"}

        elapsed_ms = int((time.perf_counter() - t0) * 1000)
        return _tabular(columns, rows[:limit], elapsed_ms, truncated=len(rows) > limit)
    finally:
        client.close()


def _run_trino(sql: str, limit: int = 200) -> dict[str, Any]:
    t0 = time.perf_counter()
    headers = {"X-Trino-User": TRINO_USER, "Content-Type": "text/plain"}
    with httpx.Client(timeout=120.0) as client:
        resp = client.post(f"{TRINO_URL}/v1/statement", content=sql, headers=headers)
        if resp.status_code >= 400:
            return {"ok": False, "error": resp.text[:500]}
        data = resp.json()
        if data.get("error"):
            return {"ok": False, "error": str(data["error"])[:500]}
        columns: list[str] = []
        rows: list[list[Any]] = []
        while True:
            if data.get("columns") and not columns:
                columns = [c["name"] for c in data["columns"]]
            if data.get("data"):
                rows.extend(data["data"])
                if len(rows) >= limit:
                    rows = rows[:limit]
                    break
            if data.get("error"):
                return {"ok": False, "error": str(data["error"])[:500]}
            nxt = data.get("nextUri")
            if not nxt:
                break
            data = client.get(nxt, headers=headers).json()
        elapsed_ms = int((time.perf_counter() - t0) * 1000)
        stats = data.get("stats") or {}
        return _tabular(columns, rows, elapsed_ms, engine_stats_ms=stats.get("elapsedTimeMillis"), truncated=len(rows) >= limit)


ENGINES = ("postgres", "mysql", "mongodb", "trino")


class SqlRequest(BaseModel):
    engine: str = Field(..., pattern="^(postgres|mysql|mongodb|trino)$")
    sql: str = Field(..., min_length=1, max_length=8000)


@router.get("/samples/{engine}")
async def get_samples(engine: str):
    if engine not in SAMPLES:
        return JSONResponse({"error": "unknown engine"}, status_code=404)
    return {"engine": engine, "samples": SAMPLES[engine], "connection": _connection_info(engine)}


@router.get("/health")
async def sql_health():
    out: dict[str, Any] = {}
    for eng in ENGINES:
        try:
            if eng == "postgres":
                _run_postgres("SELECT 1")
                out[eng] = {"ok": True, "host": DB_HOST, "user": PG_USER, "error": None}
            elif eng == "mysql":
                _run_mysql("SELECT 1")
                out[eng] = {"ok": True, "host": DB_HOST, "database": MYSQL_DB, "user": MYSQL_USER, "error": None}
            elif eng == "mongodb":
                _run_mongo("SHOW DATABASES")
                out[eng] = {"ok": True, "host": MONGO_HOST, "database": MONGO_DB, "error": None}
            else:
                r = _run_trino("SELECT 1")
                out[eng] = {"ok": bool(r.get("ok")), "url": TRINO_URL, "user": TRINO_USER, "error": r.get("error")}
        except Exception as exc:
            out[eng] = {"ok": False, "error": str(exc)[:200]}
    return out


def _connection_info(engine: str) -> dict[str, str]:
    if engine == "postgres":
        return {"host": DB_HOST, "port": str(PG_PORT), "database": PG_DB, "user": PG_USER}
    if engine == "mysql":
        return {"host": DB_HOST, "port": str(MYSQL_PORT), "database": MYSQL_DB, "user": MYSQL_USER}
    if engine == "mongodb":
        return {"host": MONGO_HOST, "port": str(MONGO_PORT), "database": MONGO_DB}
    return {"url": TRINO_URL, "user": TRINO_USER}


def _dispatch(engine: str, sql: str) -> dict[str, Any]:
    if engine == "postgres":
        return _run_postgres(sql)
    if engine == "mysql":
        return _run_mysql(sql)
    if engine == "mongodb":
        return _run_mongo(sql)
    return _run_trino(sql)


@router.post("/execute")
async def execute_sql(body: SqlRequest):
    sql = body.sql.strip().rstrip(";")
    if not sql:
        return JSONResponse({"ok": False, "error": "Empty query"}, status_code=400)
    try:
        result = _dispatch(body.engine, sql)
        if not result.get("ok"):
            return JSONResponse(result, status_code=422)
        return {**result, "engine": body.engine, "sql": sql}
    except Exception as exc:
        return JSONResponse({"ok": False, "error": str(exc)[:500]}, status_code=502)


@router.post("/benchmark")
async def benchmark():
    """Parallel analytics: PostgreSQL OLTP vs Trino distributed engine (5M sin() rows)."""
    results: dict[str, Any] = {}
    for engine, sql in BENCHMARK_SQL.items():
        try:
            results[engine] = _dispatch(engine, sql)
            results[engine]["sql"] = sql
        except Exception as exc:
            results[engine] = {"ok": False, "error": str(exc)[:300], "sql": sql}
    pg_ms = results.get("postgres", {}).get("elapsed_ms")
    tr_ms = results.get("trino", {}).get("elapsed_ms")
    faster = None
    speedup = None
    if pg_ms and tr_ms:
        faster = "trino" if tr_ms < pg_ms else "postgres"
        speedup = round(pg_ms / tr_ms, 2) if tr_ms < pg_ms else round(tr_ms / pg_ms, 2)
    return {
        "ok": True,
        "postgres": results.get("postgres"),
        "trino": results.get("trino"),
        "comparison": {
            "postgres_ms": pg_ms,
            "trino_ms": tr_ms,
            "faster": faster,
            "speedup_factor": speedup,
            "note": "5M row parallel compute — Trino distributes across workers; PostgreSQL single-node OLTP",
        },
    }