atc-agents/api/etl_offload.py

"""Autonomous ETL offload agent — source databases → S3 Parquet lake.

A background agent ("Lakehouse Loader") that, on a fixed cadence, pulls the next
small chunk of rows from every source database and writes it to the S3 object
store as a partitioned Parquet part:

    lake/<dataset>/dt=YYYY-MM-DD/part-<ts>.parquet

It progressively *backfills* the entire history in small chunks (so a 50M-row
table lands as thousands of small files) and then *tails* newly generated rows,
so the object store keeps filling and the analytics stay realtime. While
offloading it also accumulates a live federated business matrix (revenue by
region/status/channel, HR by dept, supply by type, telemetry by metric, …)
straight from the rows it actually moved — which therefore always reflects the
newest generated data, and powers the realtime Trino / Object-store dashboards.

Endpoints:
  GET  /api/etl/status     -> per-dataset progress, ingest rate, recent parts, feed
  GET  /api/etl/business   -> accumulated realtime federated business matrix
  POST /api/etl/run        -> trigger one offload cycle now
  POST /api/etl/config     -> {enabled, interval_s, chunk}
"""

from __future__ import annotations

import io
import os
import json
import random
import threading
import time
from collections import deque
from datetime import datetime, timezone
from typing import Any

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

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

DATASETS = [
    {"key": "orders", "label": "Sales orders", "engine": "PostgreSQL", "color": "#fbbf24"},
    {"key": "hr_events", "label": "HR events", "engine": "MySQL", "color": "#60a5fa"},
    {"key": "supply_events", "label": "Supply events", "engine": "MongoDB", "color": "#a78bfa"},
    {"key": "telemetry", "label": "Device telemetry", "engine": "Cassandra", "color": "#22d3ee"},
]

_lock = threading.Lock()
_state: dict[str, Any] = {
    "enabled": os.getenv("ETL_OFFLOAD_ENABLED", "1") not in ("0", "false", "False"),
    "interval_s": float(os.getenv("ETL_OFFLOAD_INTERVAL_SECONDS", "30")),
    "chunk": int(os.getenv("ETL_OFFLOAD_CHUNK", "20000")),
    "running_cycle": False,
    "started_at": None,
    "last_cycle_ts": 0.0,
    "last_cycle_rows": 0,
    "next_run_ts": 0.0,
    "cycles": 0,
    "datasets": {
        d["key"]: {"label": d["label"], "engine": d["engine"], "color": d["color"],
                   "parts": 0, "rows": 0, "bytes": 0, "backfilled": False, "cursor": None,
                   "keycol": None, "mode": None, "total_source": None, "last_ts": None,
                   "last_rows": 0, "last_key": None, "error": None}
        for d in DATASETS
    },
    "feed": deque(maxlen=60),
    "minute": deque(maxlen=60),      # (minute_epoch, rows, bytes)
    "series": deque(maxlen=48),      # per-cycle points for the realtime chart
}

# Accumulated federated business matrix, built from the rows we actually offload.
_business: dict[str, Any] = {
    "orders": {"count": 0, "revenue": 0.0, "by_region": {}, "by_status": {}, "by_channel": {}, "by_currency": {}},
    "hr": {"count": 0, "by_department": {}, "by_event": {}, "by_region": {}},
    "supply": {"count": 0, "amount": 0.0, "by_type": {}, "by_region": {}, "by_source": {}},
    "telemetry": {"count": 0, "by_metric": {}},
    "ts": deque(maxlen=48),          # {t, orders, revenue, rows}
}


# ── helpers ───────────────────────────────────────────────────────────────────
def _now() -> datetime:
    return datetime.now(timezone.utc)


def _bump(d: dict, k: Any, n: int = 1, val: float | None = None) -> None:
    if k is None or k == "":
        return
    k = str(k)
    if val is None:
        d[k] = d.get(k, 0) + n
    else:
        cur = d.get(k) or {"count": 0, "value": 0.0}
        cur["count"] += n
        cur["value"] += val
        d[k] = cur


def _feed(text: str, level: str = "info") -> None:
    _state["feed"].appendleft({"ts": _now().isoformat(), "text": text, "level": level})
    try:
        from main import add_feed
        add_feed("etl-guardian", f"[lakehouse] {text}", level)
    except Exception:
        pass


def _term(text: str, level: str = "info", phase: str = "offload") -> None:
    """Stream a line to the ETL Guardian terminal from this background thread."""
    try:
        from agent_terminal import emit_threadsafe
        emit_threadsafe("etl-guardian", text, level=level, phase=phase)
    except Exception:
        pass


def _scalar(v: Any) -> Any:
    import datetime as _dt
    import decimal
    if v is None or isinstance(v, (str, int, float, bool)):
        return v
    if isinstance(v, decimal.Decimal):
        return float(v)
    if isinstance(v, (_dt.datetime, _dt.date)):
        return v.isoformat()
    if isinstance(v, (dict, list)):
        return json.dumps(v, default=str)
    return str(v)


def _normalize(rows: list[dict]) -> list[dict]:
    keys: list[str] = []
    seen: set[str] = set()
    for r in rows:
        for k in r.keys():
            if k not in seen:
                seen.add(k)
                keys.append(k)
    return [{k: _scalar(r.get(k)) for k in keys} for r in rows]


def _write_parquet(dataset_key: str, rows: list[dict]) -> tuple[str, int]:
    import pyarrow as pa
    import pyarrow.parquet as pq
    table = pa.Table.from_pylist(_normalize(rows))
    buf = io.BytesIO()
    pq.write_table(table, buf, compression="snappy")
    body = buf.getvalue()
    now = _now()
    key = (f"lake/{dataset_key}/dt={now.strftime('%Y-%m-%d')}/"
           f"part-{int(now.timestamp() * 1000)}-{random.randint(1000, 9999)}.parquet")
    from storage_s3 import put_object_bytes
    put_object_bytes(key, body, "application/vnd.apache.parquet")
    return key, len(body)


# ── source readers: return (rows, new_cursor, done) ─────────────────────────────
def _read_orders(cursor, n):
    import psycopg2
    import psycopg2.extras
    import sql_console as s
    conn = psycopg2.connect(host=s.DB_HOST, port=s.PG_PORT, user=s.PG_USER, password=s.PG_PASS,
                            dbname=s.PG_DB, connect_timeout=8)
    try:
        conn.autocommit = True
        cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
        wm = int(cursor or 0)
        cur.execute("SELECT * FROM public.sales_orders WHERE order_id > %s ORDER BY order_id LIMIT %s", (wm, n))
        rows = [dict(r) for r in cur.fetchall()]
        new_cursor = rows[-1].get("order_id", wm) if rows else wm
        return rows, new_cursor, len(rows) < n
    finally:
        conn.close()


def _mysql_keycol(cur) -> tuple[str, str]:
    import sql_console as s
    try:
        cur.execute(
            "SELECT column_name FROM information_schema.columns WHERE table_schema=%s AND table_name='employee_events' "
            "AND extra LIKE '%%auto_increment%%' LIMIT 1", (s.MYSQL_DB,))
        r = cur.fetchone()
        if r:
            return (list(r.values())[0] if isinstance(r, dict) else r[0]), "key"
        cur.execute(
            "SELECT column_name FROM information_schema.key_column_usage WHERE table_schema=%s "
            "AND table_name='employee_events' AND constraint_name='PRIMARY' ORDER BY ordinal_position LIMIT 1",
            (s.MYSQL_DB,))
        r = cur.fetchone()
        if r:
            return (list(r.values())[0] if isinstance(r, dict) else r[0]), "key"
    except Exception:
        pass
    return "", "offset"


def _read_hr(cursor, n, st):
    import pymysql
    import sql_console as s
    conn = pymysql.connect(host=s.DB_HOST, port=s.MYSQL_PORT, user=s.MYSQL_USER, password=s.MYSQL_PASS,
                           database=s.MYSQL_DB, connect_timeout=8, cursorclass=pymysql.cursors.DictCursor)
    try:
        cur = conn.cursor()
        if not st.get("keycol") and st.get("mode") is None:
            kc, mode = _mysql_keycol(cur)
            st["keycol"] = kc
            st["mode"] = mode
        if st.get("mode") == "key" and st.get("keycol"):
            kc = st["keycol"]
            wm = int(cursor or 0)
            cur.execute(f"SELECT * FROM employee_events WHERE `{kc}` > %s ORDER BY `{kc}` LIMIT %s", (wm, n))
            rows = list(cur.fetchall())
            new_cursor = rows[-1].get(kc, wm) if rows else wm
        else:  # offset fallback
            off = int(cursor or 0)
            cur.execute("SELECT * FROM employee_events LIMIT %s OFFSET %s", (n, off))
            rows = list(cur.fetchall())
            new_cursor = off + len(rows)
        return rows, new_cursor, len(rows) < n
    finally:
        conn.close()


def _read_supply(cursor, n):
    from bson import ObjectId
    import sql_console as s
    cli = s._mongo_client()
    try:
        col = cli[s.MONGO_DB]["events"]
        q = {"_id": {"$gt": ObjectId(cursor)}} if cursor else {}
        docs = list(col.find(q).sort("_id", 1).limit(n))
        new_cursor = str(docs[-1]["_id"]) if docs else cursor
        for d in docs:
            d["_id"] = str(d["_id"])
        return docs, new_cursor, len(docs) < n
    finally:
        cli.close()


def _read_telemetry(cursor, n):
    from cassandra.query import SimpleStatement
    import sql_console as s
    cluster = s._cass_cluster()
    try:
        sess = cluster.connect()
        stmt = SimpleStatement(f"SELECT * FROM {s.CASS_KS}.device_metrics", fetch_size=n)
        kwargs = {}
        if cursor:
            try:
                kwargs["paging_state"] = bytes.fromhex(cursor)
            except Exception:
                kwargs = {}
        rs = sess.execute(stmt, **kwargs)
        rows = [dict(r._asdict()) for r in rs.current_rows]
        ps = rs.paging_state
        new_cursor = ps.hex() if ps else None
        return rows, new_cursor, new_cursor is None
    finally:
        cluster.shutdown()


# ── business matrix accumulation ────────────────────────────────────────────────
def _num(v) -> float:
    try:
        return float(v)
    except Exception:
        return 0.0


def _agg(dataset_key: str, rows: list[dict]) -> tuple[int, float]:
    """Fold a freshly-offloaded chunk into the live business matrix. Returns
    (order_rows, revenue) for the realtime time-series."""
    o_rows = 0
    o_rev = 0.0
    with _lock:
        if dataset_key == "orders":
            b = _business["orders"]
            for r in rows:
                amt = _num(r.get("amount") if r.get("amount") is not None else r.get("total_amount"))
                b["count"] += 1
                b["revenue"] += amt
                _bump(b["by_region"], r.get("region"), val=amt)
                _bump(b["by_status"], r.get("order_status") or r.get("status"))
                _bump(b["by_channel"], r.get("sales_channel") or r.get("channel"), val=amt)
                _bump(b["by_currency"], r.get("currency"))
                o_rows += 1
                o_rev += amt
        elif dataset_key == "hr_events":
            b = _business["hr"]
            for r in rows:
                b["count"] += 1
                _bump(b["by_department"], r.get("department"))
                _bump(b["by_event"], r.get("event_type") or r.get("event"))
                _bump(b["by_region"], r.get("region"))
        elif dataset_key == "supply_events":
            b = _business["supply"]
            for r in rows:
                amt = _num(r.get("amount"))
                b["count"] += 1
                b["amount"] += amt
                _bump(b["by_type"], r.get("type"), val=amt)
                _bump(b["by_region"], r.get("region"))
                _bump(b["by_source"], r.get("source"))
        elif dataset_key == "telemetry":
            b = _business["telemetry"]
            for r in rows:
                b["count"] += 1
                mt = r.get("metric_type") or r.get("metric_name")
                if mt:
                    cur = b["by_metric"].get(str(mt)) or {"count": 0, "sum": 0.0}
                    cur["count"] += 1
                    cur["sum"] += _num(r.get("metric_value"))
                    b["by_metric"][str(mt)] = cur
    return o_rows, o_rev


_READERS = {"orders": _read_orders, "hr_events": _read_hr,
            "supply_events": _read_supply, "telemetry": _read_telemetry}


def _offload_dataset(key: str) -> dict[str, Any]:
    st = _state["datasets"][key]
    chunk = int(_state["chunk"])
    res = {"rows": 0, "bytes": 0, "order_rows": 0, "revenue": 0.0}
    cur_txt = st["cursor"] if st.get("cursor") not in (None, "") else "<start>"
    _term(f"$ read {st.get('engine', key)} · {st.get('label', key)}  [cursor={cur_txt} · LIMIT {chunk}]",
          level="cmd", phase="extract")
    try:
        if key == "hr_events":
            rows, new_cursor, done = _read_hr(st["cursor"], chunk, st)
        else:
            rows, new_cursor, done = _READERS[key](st["cursor"], chunk)
    except Exception as exc:
        st["error"] = str(exc)[:160]
        _term(f"  ✗ extract failed: {str(exc)[:140]}", level="err", phase="extract")
        return res
    st["error"] = None
    if not rows:
        if done:
            st["backfilled"] = True
            if key == "telemetry":  # no global order — loop back to keep tailing
                st["cursor"] = None
        return res
    try:
        obj_key, nbytes = _write_parquet(key, rows)
    except Exception as exc:
        st["error"] = f"parquet/s3: {str(exc)[:140]}"
        _term(f"  ✗ parquet/s3 write failed: {str(exc)[:140]}", level="err", phase="load")
        return res
    _term(f"  ← pyarrow.write_table → s3://data/{obj_key}  ({len(rows)} rows · {nbytes / 1024:.1f} KB)",
          level="ok", phase="load")
    o_rows, o_rev = _agg(key, rows)
    with _lock:
        st["parts"] += 1
        st["rows"] += len(rows)
        st["bytes"] += nbytes
        st["cursor"] = new_cursor
        st["backfilled"] = bool(done)
        st["last_ts"] = _now().isoformat()
        st["last_rows"] = len(rows)
        st["last_key"] = obj_key
        if key == "telemetry" and done:
            st["cursor"] = None
    res.update({"rows": len(rows), "bytes": nbytes, "order_rows": o_rows, "revenue": o_rev})
    return res


def _source_totals() -> None:
    """Best-effort backfill targets so the dashboard can show progress %."""
    import sql_console as s
    d = _state["datasets"]
    try:
        if d["orders"]["total_source"] is None:
            d["orders"]["total_source"] = s._table_row_count("postgres", "public.sales_orders")
    except Exception:
        pass
    try:
        if d["hr_events"]["total_source"] is None:
            d["hr_events"]["total_source"] = s._table_row_count("mysql", "hr.employee_events")
    except Exception:
        pass
    try:
        if d["supply_events"]["total_source"] is None:
            d["supply_events"]["total_source"] = s._table_row_count("mongodb", "supplychain.events")
    except Exception:
        pass


def run_cycle() -> dict[str, Any]:
    if _state["running_cycle"]:
        return {"ok": True, "skipped": "cycle already running"}
    _state["running_cycle"] = True
    total = 0
    cbytes = 0
    order_rows = 0
    revenue = 0.0
    try:
        _term(f"═══ ETL offload cycle {_state['cycles'] + 1} — source DBs → S3 Parquet lake ═══",
              level="info", phase="cycle")
        _source_totals()
        for d in DATASETS:
            if not _state["enabled"]:
                break
            r = _offload_dataset(d["key"])
            total += r["rows"]
            cbytes += r["bytes"]
            order_rows += r["order_rows"]
            revenue += r["revenue"]
        now = time.time()
        _state["cycles"] += 1
        _state["last_cycle_ts"] = now
        _state["last_cycle_rows"] = total
        _state["next_run_ts"] = now + float(_state["interval_s"])
        if total:
            minute = int(now // 60) * 60
            if _state["minute"] and _state["minute"][-1][0] == minute:
                m, r0, b0 = _state["minute"][-1]
                _state["minute"][-1] = (minute, r0 + total, b0 + cbytes)
            else:
                _state["minute"].append((minute, total, cbytes))
            point = {"t": _now().strftime("%H:%M:%S"), "rows": total, "bytes": cbytes,
                     "orders": order_rows, "revenue": round(revenue, 2)}
            _state["series"].append(point)
            with _lock:
                _business["ts"].append(point)
            _term(f"═══ cycle {_state['cycles']} done — {total:,} rows · {cbytes / 1024:.1f} KB · €{int(revenue):,} revenue ═══",
                  level="ok", phase="cycle")
            _feed(f"offloaded {total:,} rows → S3 Parquet · €{int(revenue):,} (cycle {_state['cycles']})")
        else:
            _term(f"  cycle {_state['cycles']} — no new rows (all datasets caught up, tailing)",
                  level="info", phase="cycle")
    finally:
        _state["running_cycle"] = False
    return {"ok": True, "rows": total, "bytes": cbytes, "revenue": round(revenue, 2)}


def _loop() -> None:
    _state["started_at"] = _now().isoformat()
    time.sleep(12)  # let the API + sources settle
    while True:
        try:
            if _state["enabled"]:
                run_cycle()
        except Exception as exc:
            try:
                _feed(f"cycle error: {str(exc)[:160]}", level="err")
            except Exception:
                pass
        time.sleep(max(10.0, float(_state["interval_s"])))


threading.Thread(target=_loop, daemon=True, name="etl-offload").start()


# ── views ───────────────────────────────────────────────────────────────────────
def _top(d: dict, n: int = 12, value: bool = False) -> list[dict]:
    if value:
        items = sorted(d.items(), key=lambda kv: -(kv[1].get("value", 0) if isinstance(kv[1], dict) else kv[1]))
        return [{"key": k, "count": v.get("count", 0), "value": round(v.get("value", 0.0), 2)} for k, v in items[:n]]
    items = sorted(d.items(), key=lambda kv: -kv[1])
    return [{"key": k, "count": v} for k, v in items[:n]]


def status_view() -> dict[str, Any]:
    with _lock:
        datasets = []
        tot_parts = tot_rows = tot_bytes = 0
        for d in DATASETS:
            st = _state["datasets"][d["key"]]
            total_src = st.get("total_source")
            pct = None
            if total_src and total_src > 0:
                pct = min(100.0, round(100.0 * st["rows"] / total_src, 1))
            tot_parts += st["parts"]
            tot_rows += st["rows"]
            tot_bytes += st["bytes"]
            datasets.append({"key": d["key"], **{k: st[k] for k in (
                "label", "engine", "color", "parts", "rows", "bytes", "backfilled",
                "total_source", "last_ts", "last_rows", "last_key", "error")},
                "progress_pct": pct})
        now_min = int(time.time() // 60) * 60
        rate = 0
        for m, r, _b in _state["minute"]:
            if m >= now_min - 60:
                rate += r
        return {
            "ok": True,
            "enabled": _state["enabled"],
            "interval_s": _state["interval_s"],
            "chunk": _state["chunk"],
            "running_cycle": _state["running_cycle"],
            "cycles": _state["cycles"],
            "started_at": _state["started_at"],
            "last_cycle_ts": _state["last_cycle_ts"],
            "last_cycle_rows": _state["last_cycle_rows"],
            "next_run_ts": _state["next_run_ts"],
            "totals": {"parts": tot_parts, "rows": tot_rows, "bytes": tot_bytes},
            "rate_rows_per_min": rate,
            "datasets": datasets,
            "series": list(_state["series"]),
            "feed": list(_state["feed"])[:24],
        }


def business_view() -> dict[str, Any]:
    with _lock:
        o = _business["orders"]
        hr = _business["hr"]
        sup = _business["supply"]
        tel = _business["telemetry"]
        tel_metrics = sorted(
            ({"key": k, "count": v["count"], "avg": round(v["sum"] / v["count"], 2) if v["count"] else 0}
             for k, v in tel["by_metric"].items()), key=lambda x: -x["count"])[:12]
        # combine into a region matrix across datasets
        regions: dict[str, dict] = {}
        for k, v in o["by_region"].items():
            regions.setdefault(k, {})["orders"] = v.get("count", 0)
            regions[k]["revenue"] = round(v.get("value", 0.0), 2)
        for k, v in hr["by_region"].items():
            regions.setdefault(k, {})["hr_events"] = v
        for k, v in sup["by_region"].items():
            regions.setdefault(k, {})["supply_events"] = v
        region_matrix = sorted(
            ({"region": k, "orders": v.get("orders", 0), "revenue": v.get("revenue", 0),
              "hr_events": v.get("hr_events", 0), "supply_events": v.get("supply_events", 0)}
             for k, v in regions.items() if k), key=lambda r: -(r["revenue"] or 0))
        return {
            "ok": True,
            "generated_at": _now().isoformat(),
            "kpis": {
                "orders": o["count"], "revenue": round(o["revenue"], 2),
                "avg_order": round(o["revenue"] / o["count"], 2) if o["count"] else 0,
                "hr_events": hr["count"], "supply_events": sup["count"],
                "supply_amount": round(sup["amount"], 2), "telemetry": tel["count"],
                "rows_total": o["count"] + hr["count"] + sup["count"] + tel["count"],
            },
            "orders_by_region": _top(o["by_region"], 12, value=True),
            "orders_by_status": _top(o["by_status"], 8),
            "orders_by_channel": _top(o["by_channel"], 8, value=True),
            "orders_by_currency": _top(o["by_currency"], 6),
            "hr_by_department": _top(hr["by_department"], 10),
            "hr_by_event": _top(hr["by_event"], 8),
            "supply_by_type": _top(sup["by_type"], 10, value=True),
            "telemetry_by_metric": tel_metrics,
            "region_matrix": region_matrix,
            "ts": list(_business["ts"]),
        }


@router.get("/status")
async def get_status() -> JSONResponse:
    return JSONResponse(status_view())


@router.get("/business")
async def get_business() -> JSONResponse:
    return JSONResponse(business_view())


@router.post("/run")
async def post_run() -> JSONResponse:
    from starlette.concurrency import run_in_threadpool
    res = await run_in_threadpool(run_cycle)
    return JSONResponse({**res, "status": status_view()})


@router.post("/config")
async def post_config(body: dict = Body(default={})) -> JSONResponse:
    if "enabled" in body:
        _state["enabled"] = bool(body["enabled"])
    if "interval_s" in body:
        try:
            _state["interval_s"] = max(10.0, min(900.0, float(body["interval_s"])))
        except Exception:
            pass
    if "chunk" in body:
        try:
            _state["chunk"] = max(200, min(50000, int(body["chunk"])))
        except Exception:
            pass
    _feed(f"config updated · interval={_state['interval_s']}s · chunk={_state['chunk']} · enabled={_state['enabled']}")
    return JSONResponse({"ok": True, "enabled": _state["enabled"],
                         "interval_s": _state["interval_s"], "chunk": _state["chunk"]})