atc-agents/api/sql_console.py

"""Live SQL console — PostgreSQL + Trino with benchmark."""

from __future__ import annotations

import os
import time
from typing import Any

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

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

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');"},
    ],
    "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 _run_postgres(sql: str, limit: int = 200) -> dict[str, Any]:
    t0 = time.perf_counter()
    conn = psycopg2.connect(
        host=PG_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 = cur.fetchmany(limit)
            data = [list(r) for r in rows]
            return {"ok": True, "columns": columns, "rows": data, "row_count": len(data), "elapsed_ms": elapsed_ms, "truncated": len(data) >= limit}
        return {"ok": True, "columns": [], "rows": [], "row_count": 0, "elapsed_ms": elapsed_ms, "message": "OK"}
    finally:
        conn.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 {
            "ok": True,
            "columns": columns,
            "rows": rows,
            "row_count": len(rows),
            "elapsed_ms": elapsed_ms,
            "engine_stats_ms": stats.get("elapsedTimeMillis"),
            "truncated": len(rows) >= limit,
        }


class SqlRequest(BaseModel):
    engine: str = Field(..., pattern="^(postgres|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():
    pg_ok = trino_ok = False
    pg_err = trino_err = None
    try:
        _run_postgres("SELECT 1")
        pg_ok = True
    except Exception as exc:
        pg_err = str(exc)[:200]
    try:
        r = _run_trino("SELECT 1")
        trino_ok = bool(r.get("ok"))
        if not trino_ok:
            trino_err = r.get("error")
    except Exception as exc:
        trino_err = str(exc)[:200]
    return {
        "postgres": {"ok": pg_ok, "host": PG_HOST, "user": PG_USER, "error": pg_err},
        "trino": {"ok": trino_ok, "url": TRINO_URL, "user": TRINO_USER, "error": trino_err},
    }


def _connection_info(engine: str) -> dict[str, str]:
    if engine == "postgres":
        return {"host": PG_HOST, "port": str(PG_PORT), "database": PG_DB, "user": PG_USER}
    return {"url": TRINO_URL, "user": TRINO_USER}


@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:
        if body.engine == "postgres":
            result = _run_postgres(sql)
        else:
            result = _run_trino(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:
            if engine == "postgres":
                results[engine] = _run_postgres(sql)
            else:
                results[engine] = _run_trino(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
    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 and tr_ms < pg_ms else round(tr_ms / pg_ms, 2) if pg_ms else None
    else:
        speedup = None
    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",
        },
    }