atc-agents/api/spark_workbench.py

"""Databricks-style lakehouse workbench for the Command Center.

Lets the user interactively select data from any federated source (Iceberg,
Hive/HDFS, Postgres, MySQL, Mongo, Cassandra, Kafka) and run distributed
transformations on the lakehouse compute layer:

    explore -> transform (aggregate / filter / join / profile) -> materialize to Iceberg (S3/HDFS)

Execution streams live engine metrics (state, splits, rows, bytes, CPU, wall,
peak memory, nodes) so the UI can show a live "what the cluster is doing" matrix
exactly like Databricks' Spark UI. Compute runs on the lakehouse engine (Trino
coordinator on the Spark cluster host) which distributes work across the
workers; Spark batch DAGs remain available via the jobs API.
"""

from __future__ import annotations

import asyncio
import os
import re
import time
import uuid
from typing import Any

import httpx
from fastapi import APIRouter, Body, Query
from fastapi.responses import JSONResponse

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

router = APIRouter(prefix="/api/spark", tags=["spark-workbench"])

_runs: dict[str, dict[str, Any]] = {}
_run_order: list[str] = []
_MAX_RUNS = 40
_MAX_RESULT_ROWS = 500

_IDENT_RE = re.compile(r"^[A-Za-z0-9_.\" ]+$")
_TABLE_RE = re.compile(r"^[A-Za-z0-9_.\"]+$")

AGG_FUNCS = {"count", "sum", "avg", "min", "max", "approx_distinct", "count_distinct"}
JOIN_TYPES = {"INNER", "LEFT", "RIGHT", "FULL"}


def _feed(agent_id: str, message: str, level: str = "info") -> None:
    try:
        from main import add_feed
        add_feed(agent_id, message, level)
    except Exception:
        pass


def _qident(name: str) -> str:
    name = name.strip().strip('"')
    if not re.match(r"^[A-Za-z0-9_]+$", name):
        raise ValueError(f"invalid identifier: {name}")
    return f'"{name}"'


def _safe_table(name: str) -> str:
    name = (name or "").strip()
    if not name or not _TABLE_RE.match(name):
        raise ValueError(f"invalid table reference: {name}")
    return name


# ── Trino helpers ────────────────────────────────────────────────────────────
def _trino_headers() -> dict[str, str]:
    return {"X-Trino-User": TRINO_USER, "Content-Type": "text/plain"}


async def _trino_collect(sql: str, limit: int = 200, timeout: float = 60.0) -> dict[str, Any]:
    """Synchronous-style helper: run a query and return all rows (for catalog ops)."""
    columns: list[str] = []
    rows: list[list[Any]] = []
    async with httpx.AsyncClient(timeout=timeout) as client:
        r = await client.post(f"{TRINO_URL}/v1/statement", content=sql.encode(), headers=_trino_headers())
        if r.status_code >= 400:
            return {"ok": False, "error": r.text[:400]}
        data = r.json()
        while True:
            if data.get("error"):
                return {"ok": False, "error": str(data["error"])[:400]}
            if data.get("columns") and not columns:
                columns = [c["name"] for c in data["columns"]]
            for row in data.get("data") or []:
                rows.append(row)
                if len(rows) >= limit:
                    break
            nxt = data.get("nextUri")
            if not nxt or len(rows) >= limit:
                if nxt:
                    try:
                        await client.delete(nxt, headers=_trino_headers())
                    except Exception:
                        pass
                break
            data = (await client.get(nxt, headers=_trino_headers())).json()
    return {"ok": True, "columns": columns, "rows": rows}


def _norm_stats(stats: dict[str, Any]) -> dict[str, Any]:
    return {
        "state": stats.get("state"),
        "nodes": stats.get("nodes"),
        "total_splits": stats.get("totalSplits"),
        "queued_splits": stats.get("queuedSplits"),
        "running_splits": stats.get("runningSplits"),
        "completed_splits": stats.get("completedSplits"),
        "processed_rows": stats.get("processedRows"),
        "processed_bytes": stats.get("processedBytes"),
        "physical_input_bytes": stats.get("physicalInputBytes"),
        "peak_memory_bytes": stats.get("peakMemoryBytes"),
        "cpu_time_ms": stats.get("cpuTimeMillis"),
        "wall_time_ms": stats.get("wallTimeMillis"),
        "elapsed_ms": stats.get("elapsedTimeMillis"),
        "progress_pct": round(
            (stats.get("completedSplits") or 0) / stats["totalSplits"] * 100, 1
        ) if stats.get("totalSplits") else (100.0 if stats.get("state") == "FINISHED" else 0.0),
    }


async def _execute_run(run_id: str, sql: str, returns_rows: bool, pre_sql: str | None = None) -> None:
    run = _runs[run_id]
    run["state"] = "RUNNING"
    columns: list[str] = []
    rows: list[list[Any]] = []
    try:
        if pre_sql:
            pre = await _trino_collect(pre_sql, 1)
            if not pre.get("ok"):
                run.update(state="FAILED", error=pre.get("error"), ended_at=time.time())
                _feed("lakehouse-ops", f"[workbench] {run['label']}: pre-step failed", "err")
                return
        async with httpx.AsyncClient(timeout=None) as client:
            r = await client.post(f"{TRINO_URL}/v1/statement", content=sql.encode(), headers=_trino_headers())
            if r.status_code >= 400:
                run.update(state="FAILED", error=r.text[:500], ended_at=time.time())
                _feed("lakehouse-ops", f"[workbench] {run['label']}: failed ({r.status_code})", "err")
                return
            data = r.json()
            run["query_id"] = data.get("id")
            while True:
                if run.get("cancel_requested"):
                    nxt = data.get("nextUri")
                    if nxt:
                        try:
                            await client.delete(nxt, headers=_trino_headers())
                        except Exception:
                            pass
                    run.update(state="CANCELED", ended_at=time.time())
                    _feed("lakehouse-ops", f"[workbench] {run['label']}: canceled", "warn")
                    return
                if data.get("stats"):
                    run["stats"] = _norm_stats(data["stats"])
                    st = data["stats"].get("state")
                    if st:
                        run["engine_state"] = st
                if data.get("error"):
                    run.update(state="FAILED", error=str(data["error"])[:500], ended_at=time.time())
                    _feed("lakehouse-ops", f"[workbench] {run['label']}: {str(data['error'])[:120]}", "err")
                    return
                if data.get("columns") and not columns:
                    columns = [c["name"] for c in data["columns"]]
                    run["columns"] = columns
                for row in data.get("data") or []:
                    if returns_rows and len(rows) < _MAX_RESULT_ROWS:
                        rows.append(row)
                if data.get("updateType"):
                    run["update_type"] = data.get("updateType")
                nxt = data.get("nextUri")
                if not nxt:
                    break
                data = (await client.get(nxt, headers=_trino_headers())).json()
                run["next_uri"] = nxt
        run["rows"] = rows
        run["row_count"] = len(rows)
        run.update(state="FINISHED", ended_at=time.time())
        if run.get("stats"):
            run["stats"]["state"] = "FINISHED"
            run["stats"]["progress_pct"] = 100.0
        msg = f"[workbench] {run['label']}: finished"
        if run.get("target"):
            msg = f"[workbench] {run['label']}: materialized → {run['target']}"
        _feed("lakehouse-ops", msg, "info")
    except Exception as exc:
        run.update(state="FAILED", error=str(exc)[:500], ended_at=time.time())
        _feed("lakehouse-ops", f"[workbench] {run['label']}: error {str(exc)[:120]}", "err")


# ── SQL builders ─────────────────────────────────────────────────────────────
def _build_sql(body: dict[str, Any]) -> tuple[str, bool, str, str | None, str | None]:
    """Return (sql, returns_rows, label, materialize_target, pre_sql)."""
    op = body.get("operation", "preview")
    limit = max(1, min(int(body.get("limit", 200)), _MAX_RESULT_ROWS))
    materialize = body.get("materialize") or {}
    target = None

    if op == "sql":
        sql = (body.get("sql") or "").strip().rstrip(";")
        if not sql:
            raise ValueError("empty SQL")
        select_sql = sql
        label = "Custom SQL"
        returns_rows = sql.lower().lstrip().startswith(("select", "show", "describe", "with", "explain"))

    elif op == "preview":
        table = _safe_table(body.get("table"))
        select_sql = f"SELECT * FROM {table} LIMIT {limit}"
        label = f"Preview {table}"
        returns_rows = True

    elif op == "filter":
        table = _safe_table(body.get("table"))
        where = (body.get("where") or "").strip()
        clause = f" WHERE {where}" if where else ""
        select_sql = f"SELECT * FROM {table}{clause} LIMIT {limit}"
        label = f"Filter {table}"
        returns_rows = True

    elif op == "aggregate":
        table = _safe_table(body.get("table"))
        group_by = [c for c in (body.get("group_by") or []) if c]
        metrics = body.get("metrics") or []
        select_parts: list[str] = [_qident(c) for c in group_by]
        for m in metrics:
            fn = (m.get("fn") or "count").lower()
            if fn not in AGG_FUNCS:
                raise ValueError(f"unsupported function {fn}")
            col = m.get("col")
            alias = m.get("alias") or (f"{fn}_{col}" if col else fn)
            if fn == "count" and (not col or col == "*"):
                expr = "count(*)"
            elif fn == "count_distinct":
                expr = f"count(DISTINCT {_qident(col)})"
            else:
                expr = f"{fn}({_qident(col)})"
            select_parts.append(f"{expr} AS {_qident(alias)}")
        if not select_parts:
            select_parts = ["count(*) AS cnt"]
        gb = f" GROUP BY {', '.join(_qident(c) for c in group_by)}" if group_by else ""
        order = ""
        if group_by:
            order = f" ORDER BY {', '.join(_qident(c) for c in group_by)}"
        select_sql = f"SELECT {', '.join(select_parts)} FROM {table}{gb}{order} LIMIT {limit}"
        label = f"Aggregate {table}"
        returns_rows = True

    elif op == "join":
        left = _safe_table(body.get("left"))
        right = _safe_table(body.get("right"))
        jt = (body.get("join_type") or "INNER").upper()
        if jt not in JOIN_TYPES:
            raise ValueError(f"invalid join type {jt}")
        lk = _qident(body.get("left_key"))
        rk = _qident(body.get("right_key"))
        select_sql = (
            f"SELECT l.*, r.* FROM {left} l {jt} JOIN {right} r "
            f"ON l.{lk} = r.{rk} LIMIT {limit}"
        )
        label = f"Join {left} ⋈ {right}"
        returns_rows = True

    elif op == "profile":
        table = _safe_table(body.get("table"))
        cols = [c for c in (body.get("columns") or []) if c][:12]
        parts = ["count(*) AS row_count"]
        for c in cols:
            qc = _qident(c)
            parts.append(f"approx_distinct({qc}) AS {_qident(c + '_distinct')}")
            parts.append(f"count({qc}) AS {_qident(c + '_nonnull')}")
        select_sql = f"SELECT {', '.join(parts)} FROM {table}"
        label = f"Profile {table}"
        returns_rows = True

    else:
        raise ValueError(f"unknown operation {op}")

    if materialize.get("enabled"):
        schema = materialize.get("schema", "hadoop")
        name = materialize.get("table")
        if not name:
            raise ValueError("materialize target table required")
        target = f"iceberg.{_qident(schema).strip(chr(34))}.{_qident(name).strip(chr(34))}"
        mode = (materialize.get("mode") or "create").lower()
        pre = None
        if mode == "replace":
            pre = f"DROP TABLE IF EXISTS {target}"
            ddl = f"CREATE TABLE {target} AS {select_sql}"
        elif mode == "insert":
            ddl = f"INSERT INTO {target} {select_sql}"
        else:
            ddl = f"CREATE TABLE {target} AS {select_sql}"
        return ddl, False, f"Materialize → {target}", target, pre

    return select_sql, returns_rows, label, None, None


# ── catalog endpoints ────────────────────────────────────────────────────────
@router.get("/catalogs")
async def list_catalogs() -> JSONResponse:
    res = await _trino_collect("SHOW CATALOGS", 100)
    if not res.get("ok"):
        return JSONResponse(res, status_code=502)
    cats = [r[0] for r in res["rows"] if r[0] not in ("system",)]
    return JSONResponse({"ok": True, "catalogs": cats})


@router.get("/schemas")
async def list_schemas(catalog: str = Query(...)) -> JSONResponse:
    cat = _safe_table(catalog)
    res = await _trino_collect(f"SHOW SCHEMAS FROM {cat}", 200)
    if not res.get("ok"):
        return JSONResponse(res, status_code=502)
    skip = {"information_schema"}
    schemas = [r[0] for r in res["rows"] if r[0] not in skip]
    return JSONResponse({"ok": True, "catalog": catalog, "schemas": schemas})


@router.get("/tables")
async def list_tables(catalog: str = Query(...), schema: str = Query(...)) -> JSONResponse:
    cat = _safe_table(catalog)
    sch = _safe_table(schema)
    res = await _trino_collect(f"SHOW TABLES FROM {cat}.{sch}", 500)
    if not res.get("ok"):
        return JSONResponse(res, status_code=502)
    tables = [{"name": r[0], "fqn": f"{catalog}.{schema}.{r[0]}"} for r in res["rows"]]
    return JSONResponse({"ok": True, "catalog": catalog, "schema": schema, "tables": tables})


@router.get("/columns")
async def list_columns(table: str = Query(...)) -> JSONResponse:
    tbl = _safe_table(table)
    res = await _trino_collect(f"DESCRIBE {tbl}", 500)
    if not res.get("ok"):
        return JSONResponse(res, status_code=502)
    cols = [{"name": r[0], "type": r[1] if len(r) > 1 else ""} for r in res["rows"]]
    return JSONResponse({"ok": True, "table": table, "columns": cols})


# ── run endpoints ────────────────────────────────────────────────────────────
@router.post("/run")
async def create_run(body: dict[str, Any] = Body(...)) -> JSONResponse:
    try:
        sql, returns_rows, label, target, pre_sql = _build_sql(body)
    except ValueError as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=400)

    run_id = uuid.uuid4().hex[:12]
    _runs[run_id] = {
        "id": run_id,
        "operation": body.get("operation", "preview"),
        "label": label,
        "sql": sql,
        "target": target,
        "state": "QUEUED",
        "engine_state": "QUEUED",
        "stats": {},
        "columns": [],
        "rows": [],
        "row_count": None,
        "error": None,
        "started_at": time.time(),
        "ended_at": None,
        "cancel_requested": False,
    }
    _run_order.append(run_id)
    while len(_run_order) > _MAX_RUNS:
        old = _run_order.pop(0)
        _runs.pop(old, None)

    _feed("lakehouse-ops", f"[workbench] {label}: submitted", "info")
    asyncio.create_task(_execute_run(run_id, sql, returns_rows, pre_sql))
    return JSONResponse({"ok": True, "run_id": run_id, "sql": sql, "label": label, "target": target})


def _run_public(run: dict[str, Any], include_rows: bool = True) -> dict[str, Any]:
    out = {k: v for k, v in run.items() if k not in ("next_uri", "cancel_requested")}
    if not include_rows:
        out.pop("rows", None)
    return out


@router.get("/run/{run_id}")
async def get_run(run_id: str) -> JSONResponse:
    run = _runs.get(run_id)
    if not run:
        return JSONResponse({"ok": False, "error": "unknown run"}, status_code=404)
    return JSONResponse({"ok": True, "run": _run_public(run)})


@router.post("/run/{run_id}/cancel")
async def cancel_run(run_id: str) -> JSONResponse:
    run = _runs.get(run_id)
    if not run:
        return JSONResponse({"ok": False, "error": "unknown run"}, status_code=404)
    run["cancel_requested"] = True
    return JSONResponse({"ok": True, "run_id": run_id, "state": "canceling"})


@router.get("/runs")
async def list_runs() -> JSONResponse:
    out = [_run_public(_runs[r], include_rows=False) for r in reversed(_run_order) if r in _runs]
    return JSONResponse({"ok": True, "runs": out})


@router.get("/live")
async def spark_live() -> JSONResponse:
    """Live cluster + active-run matrix for the workbench dashboard."""
    spark: dict[str, Any] = {}
    try:
        from streaming_ops import collect_spark
        spark = await collect_spark()
    except Exception as exc:
        spark = {"error": str(exc)[:200]}

    active = [
        _run_public(_runs[r], include_rows=False)
        for r in reversed(_run_order)
        if r in _runs and _runs[r]["state"] in ("QUEUED", "RUNNING")
    ]
    recent = [
        _run_public(_runs[r], include_rows=False)
        for r in reversed(_run_order[-8:])
        if r in _runs
    ]
    return JSONResponse({
        "ok": True,
        "spark": spark,
        "active_runs": active,
        "recent_runs": recent,
        "ts": time.time(),
    })