"""Live Spark + Kafka visibility and manual job control for the Command Center.""" from __future__ import annotations import asyncio import os import time from typing import Any import httpx from fastapi import APIRouter, Body from fastapi.responses import JSONResponse, Response SPARK_UI_URL = os.getenv("SPARK_UI_URL", "http://10.0.21.50:8080").rstrip("/") KAFKA_UI_URL = os.getenv("KAFKA_UI_URL", "http://10.0.21.36:9000").rstrip("/") KAFKA_CONNECT_URL = os.getenv("KAFKA_CONNECT_URL", "http://10.0.21.50:8083").rstrip("/") AIRFLOW_URL = os.getenv("AIRFLOW_URL", "http://10.0.21.55:8080").rstrip("/") AIRFLOW_USER = os.getenv("AIRFLOW_USER", "admin") AIRFLOW_PASSWORD = os.getenv("AIRFLOW_PASSWORD", "") from hdfs_kafka import export_hdfs_to_kafka, hdfs_export_snapshot router = APIRouter(prefix="/api/pipeline/streaming", tags=["streaming"]) # Manual Spark / lakehouse jobs (Airflow DAGs on the lab) SPARK_JOBS: dict[str, dict[str, Any]] = { "spark_to_curated": { "label": "Spark → Iceberg curated (mask PII)", "dag_id": "mask_to_curated", "agent": "lakehouse-ops", "description": "Runs the mask_to_curated DAG — Spark/SQL transform into iceberg.curated_masked", "default_conf": {}, "editable_fields": ["batch_size", "sources"], }, "hadoop_to_trino": { "label": "HDFS → Iceberg (historical)", "dag_id": "hadoop_to_trino", "agent": "hadoop-ranger", "description": "Load historical_sales from HDFS into Iceberg via Trino/Spark pipeline", "default_conf": {"mode": "refresh"}, "editable_fields": ["mode"], }, "spark_to_s3": { "label": "Spark → S3 curated layer", "dag_id": "mask_to_curated", "agent": "lakehouse-ops", "description": "Spark transform into Iceberg curated tables mirrored to S3", "default_conf": {"target": "s3"}, "editable_fields": ["target", "batch_size"], }, "generate_all": { "label": "Seed all source DBs", "dag_id": "generate_data_all_databases", "agent": "data-custodian", "description": "Airflow DAG — generates rows into PostgreSQL, MySQL, MongoDB, Cassandra, Neo4j", "default_conf": {"rows": 3000}, "editable_fields": ["rows"], }, } # Master pipeline pulse switch: running | paused | stopped _flow_state: dict[str, Any] = {"mode": "running", "since": time.time()} def flow_mode() -> str: return _flow_state.get("mode", "running") def flow_snapshot() -> dict[str, Any]: return {"mode": _flow_state.get("mode", "running"), "since": _flow_state.get("since")} _cache: dict[str, Any] = {"ts": 0.0, "data": None} _TTL = 8.0 _token_cache: dict[str, Any] = {"token": None, "exp": 0.0} 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 async def _airflow_token(client: httpx.AsyncClient) -> str: now = time.time() if _token_cache["token"] and _token_cache["exp"] > now + 30: return _token_cache["token"] r = await client.post( f"{AIRFLOW_URL}/auth/token", json={"username": AIRFLOW_USER, "password": AIRFLOW_PASSWORD}, timeout=10, ) r.raise_for_status() tok = r.json()["access_token"] _token_cache["token"] = tok _token_cache["exp"] = now + 20 * 60 return tok async def collect_spark() -> dict[str, Any]: out: dict[str, Any] = { "ui_url": SPARK_UI_URL, "ui_ok": False, "status": "UNKNOWN", "workers": [], "alive_workers": 0, "cores": 0, "cores_used": 0, "memory_mb": 0, "memory_used_mb": 0, "active_apps": [], "completed_apps": [], } try: async with httpx.AsyncClient(timeout=6.0) as client: r = await client.get(f"{SPARK_UI_URL}/json/") if r.status_code >= 400: return out d = r.json() out["ui_ok"] = True out["status"] = d.get("status") or "ALIVE" out["alive_workers"] = int(d.get("aliveworkers") or 0) out["cores"] = int(d.get("cores") or 0) out["cores_used"] = int(d.get("coresused") or 0) out["memory_mb"] = int(d.get("memory") or 0) out["memory_used_mb"] = int(d.get("memoryused") or 0) workers = d.get("workers") or [] out["workers"] = [ { "id": w.get("id"), "host": w.get("host"), "cores": w.get("cores"), "cores_used": w.get("coresused"), "memory_mb": w.get("memory"), "state": w.get("state"), "webui": w.get("webuiaddress"), } for w in workers ] for app in (d.get("activeapps") or [])[:20]: out["active_apps"].append({ "id": app.get("id"), "name": app.get("name"), "cores": app.get("cores"), "memory_mb": app.get("memory"), "submitdate": app.get("submitdate"), "duration_ms": app.get("duration"), "user": app.get("user"), }) for app in (d.get("completedapps") or [])[:10]: out["completed_apps"].append({ "id": app.get("id"), "name": app.get("name"), "duration_ms": app.get("duration"), }) except Exception as exc: out["error"] = str(exc)[:200] return out async def collect_kafka() -> dict[str, Any]: out: dict[str, Any] = { "ui_url": KAFKA_UI_URL, "ui_ok": False, "connect_url": KAFKA_CONNECT_URL, "connect_ok": False, "cluster": {}, "topics": [], "connectors": [], } try: async with httpx.AsyncClient(timeout=8.0) as client: cr = await client.get(f"{KAFKA_UI_URL}/api/clusters") if cr.status_code < 400: clusters = cr.json() out["ui_ok"] = True if clusters: name = clusters[0].get("name", "local") out["cluster"] = { "name": name, "status": clusters[0].get("status"), "broker_count": clusters[0].get("brokerCount"), "topic_count": clusters[0].get("topicCount"), "online_partitions": clusters[0].get("onlinePartitionCount"), } tr = await client.get( f"{KAFKA_UI_URL}/api/clusters/{name}/topics", params={"page": 1, "perPage": 50, "showInternal": False}, ) if tr.status_code < 400: topics = tr.json().get("topics") or [] out["topics"] = [ { "name": t.get("name"), "partitions": t.get("partitionCount"), "replicas": t.get("replicationFactor"), "messages": t.get("messagesCount"), } for t in topics[:50] ] lr = await client.get(f"{KAFKA_CONNECT_URL}/connectors") if lr.status_code < 400: out["connect_ok"] = True names_raw = lr.json() names = names_raw if isinstance(names_raw, list) else [] for cn in names[:20]: try: sr = await client.get(f"{KAFKA_CONNECT_URL}/connectors/{cn}/status") st = sr.json() if sr.status_code < 400 else {} conn = st.get("connector") or {} tasks = st.get("tasks") or [] out["connectors"].append({ "name": cn, "state": conn.get("state"), "worker": conn.get("worker_id"), "tasks": [{"id": t.get("id"), "state": t.get("state")} for t in tasks], "type": st.get("type"), }) except Exception: out["connectors"].append({"name": cn, "state": "UNKNOWN"}) except Exception as exc: out["error"] = str(exc)[:200] return out async def build_streaming_status() -> dict[str, Any]: spark, kafka = await asyncio.gather(collect_spark(), collect_kafka()) cdc_active = False try: from cdc_consumer import snapshot as cdc_snapshot cdc = cdc_snapshot(15) cdc_active = bool(cdc.get("connected")) and int(cdc.get("window_total") or 0) > 0 except Exception: pass connectors_running = sum( 1 for c in kafka.get("connectors", []) if (c.get("state") or "").upper() == "RUNNING" ) spark_alive = spark.get("ui_ok") and (spark.get("status") or "").upper() == "ALIVE" apps_running = len(spark.get("active_apps") or []) hdfs_recent = hdfs_export_snapshot().get("recent") edges = { "hdfs→kafka": bool(hdfs_recent), "kafka→spark": cdc_active and spark_alive or hdfs_recent, "spark→iceberg": apps_running > 0 or spark_alive, "spark→s3": apps_running > 0 or spark_alive, "sources→kafka": cdc_active, "connectors": connectors_running > 0, } mode = flow_mode() if mode != "running": edges = {k: False for k in edges} return { "ok": True, "flow": mode, "spark": spark, "kafka": kafka, "edges": edges, "jobs": [ {**{"id": k}, **{kk: vv for kk, vv in v.items() if kk != "editable_fields"}} for k, v in SPARK_JOBS.items() ], "ts": time.time(), } @router.get("/status") async def streaming_status(refresh: bool = False) -> JSONResponse: now = time.time() if not refresh and _cache["data"] and now - _cache["ts"] < _TTL: return JSONResponse(_cache["data"]) data = await build_streaming_status() _cache["data"] = data _cache["ts"] = now return JSONResponse(data) @router.get("/jobs") async def list_spark_jobs() -> JSONResponse: return JSONResponse({"ok": True, "jobs": SPARK_JOBS}) @router.post("/flow/{action}") async def set_flow(action: str) -> JSONResponse: mapping = { "pause": "paused", "stop": "stopped", "resume": "running", "start": "running", "run": "running", } if action not in mapping: return JSONResponse({"ok": False, "error": f"unknown action {action}"}, status_code=400) _flow_state["mode"] = mapping[action] _flow_state["since"] = time.time() _cache["ts"] = 0 label = {"running": "resumed", "paused": "paused", "stopped": "stopped"}[mapping[action]] _feed("lakehouse-ops", f"[pipeline] data flow {label}", "info" if mapping[action] == "running" else "warn") return JSONResponse({"ok": True, "flow": flow_snapshot()}) @router.get("/flow") async def get_flow() -> JSONResponse: return JSONResponse({"ok": True, "flow": flow_snapshot()}) @router.post("/jobs/{job_id}/trigger") async def trigger_spark_job(job_id: str, body: dict[str, Any] = Body(default={})) -> JSONResponse: job = SPARK_JOBS.get(job_id) if not job: return JSONResponse({"ok": False, "error": f"unknown job {job_id}"}, status_code=400) conf = {**(job.get("default_conf") or {}), **(body.get("conf") or {})} agent = job.get("agent", "lakehouse-ops") dag_id = job["dag_id"] _feed(agent, f"[spark] Triggering {job['label']} conf={conf}", "info") try: async with httpx.AsyncClient() as client: tok = await _airflow_token(client) h = {"Authorization": f"Bearer {tok}"} r = await client.post( f"{AIRFLOW_URL}/api/v2/dags/{dag_id}/dagRuns", headers=h, json={"logical_date": None, "conf": conf}, timeout=20, ) if r.status_code >= 400: _feed(agent, f"[spark] {job['label']}: Airflow {r.status_code}", "err") return JSONResponse({"ok": False, "error": f"airflow {r.status_code}: {r.text[:300]}"}, status_code=502) run = r.json() run_id = run.get("dag_run_id") _feed(agent, f"[spark] {job['label']}: started run {run_id}", "info") return JSONResponse({"ok": True, "job_id": job_id, "dag_id": dag_id, "run_id": run_id, "conf": conf}) except Exception as exc: _feed(agent, f"[spark] {job['label']}: {str(exc)[:120]}", "err") return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) @router.post("/hdfs/to-kafka") async def hdfs_to_kafka(body: dict[str, Any] = Body(default={})) -> JSONResponse: path = body.get("path", "/data/historical/sales_orders") topic = body.get("topic", "hdfs.historical.sales") limit = int(body.get("limit", 2000)) try: result = await export_hdfs_to_kafka(path=path, topic=topic, limit=limit, feed=_feed) if not result.get("ok"): return JSONResponse(result, status_code=502) _cache["ts"] = 0 return JSONResponse(result) except Exception as exc: _feed("hadoop-ranger", f"[hdfs→kafka] failed: {str(exc)[:120]}", "err") return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) @router.post("/pipeline/{pipeline_id}") async def run_streaming_pipeline(pipeline_id: str) -> JSONResponse: if pipeline_id != "hadoop-lake": return JSONResponse({"ok": False, "error": f"unknown pipeline {pipeline_id}"}, status_code=400) steps: list[str] = [] try: exp = await export_hdfs_to_kafka(feed=_feed) if not exp.get("ok"): return JSONResponse({"ok": False, "error": exp.get("error"), "steps": steps}, status_code=502) steps.append(f"hdfs→kafka ({exp.get('rows_sent')} rows)") async with httpx.AsyncClient() as client: tok = await _airflow_token(client) h = {"Authorization": f"Bearer {tok}"} for job_id, dag_id, agent in [ ("hadoop_to_trino", "hadoop_to_trino", "hadoop-ranger"), ("spark_to_curated", "mask_to_curated", "lakehouse-ops"), ]: r = await client.post( f"{AIRFLOW_URL}/api/v2/dags/{dag_id}/dagRuns", headers=h, json={"logical_date": None, "conf": {}}, timeout=20, ) if r.status_code >= 400: _feed(agent, f"[pipeline] {dag_id} failed {r.status_code}", "err") return JSONResponse({"ok": False, "error": f"{dag_id}: {r.text[:200]}", "steps": steps}, status_code=502) steps.append(dag_id) _feed(agent, f"[pipeline] started {dag_id}", "info") _cache["ts"] = 0 return JSONResponse({"ok": True, "pipeline": pipeline_id, "steps": steps, "export": exp}) except Exception as exc: return JSONResponse({"ok": False, "error": str(exc), "steps": steps}, status_code=500) @router.post("/kafka/connectors/{name}/restart") async def restart_kafka_connector(name: str) -> JSONResponse: _feed("etl-guardian", f"[kafka] Restarting connector {name}", "info") try: async with httpx.AsyncClient(timeout=15.0) as client: r = await client.post(f"{KAFKA_CONNECT_URL}/connectors/{name}/restart") if r.status_code >= 400: return JSONResponse({"ok": False, "error": r.text[:300]}, status_code=502) _feed("etl-guardian", f"[kafka] Connector {name} restart requested", "info") _cache["ts"] = 0 return JSONResponse({"ok": True, "connector": name, "action": "restart"}) except Exception as exc: return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) @router.post("/kafka/connectors/{name}/pause") async def pause_kafka_connector(name: str) -> JSONResponse: try: async with httpx.AsyncClient(timeout=10.0) as client: r = await client.put(f"{KAFKA_CONNECT_URL}/connectors/{name}/pause") if r.status_code >= 400: return JSONResponse({"ok": False, "error": r.text[:300]}, status_code=502) _cache["ts"] = 0 return JSONResponse({"ok": True, "connector": name, "action": "pause"}) except Exception as exc: return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) @router.post("/kafka/connectors/{name}/resume") async def resume_kafka_connector(name: str) -> JSONResponse: try: async with httpx.AsyncClient(timeout=10.0) as client: r = await client.put(f"{KAFKA_CONNECT_URL}/connectors/{name}/resume") if r.status_code >= 400: return JSONResponse({"ok": False, "error": r.text[:300]}, status_code=502) _cache["ts"] = 0 return JSONResponse({"ok": True, "connector": name, "action": "resume"}) except Exception as exc: return JSONResponse({"ok": False, "error": str(exc)}, status_code=500)