From bac3a907750180f15e122e67802072487d825331 Mon Sep 17 00:00:00 2001 From: mo Date: Fri, 26 Jun 2026 01:12:07 +0000 Subject: [PATCH] api: pipeline_ops router (per-source generate trigger, run status, Trino sync counts) --- api/Dockerfile | 2 +- api/main.py | 2 + api/pipeline_ops.py | 165 ++++++++++++++++++++++++++++++++++++++++++++ docker-compose.yml | 2 + 4 files changed, 170 insertions(+), 1 deletion(-) create mode 100644 api/pipeline_ops.py diff --git a/api/Dockerfile b/api/Dockerfile index 143d313..db17ed5 100644 --- a/api/Dockerfile +++ b/api/Dockerfile @@ -4,7 +4,7 @@ WORKDIR /app RUN apt-get update && apt-get install -y --no-install-recommends curl && rm -rf /var/lib/apt/lists/* COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt -COPY main.py lab_context.py agent_terminal.py workload.py node_registry.py node_ops.py topology_views.py supervisor.py approval_service.py db.py dockhand_envs.py presentation.py database_inventory.py presentation_upload.py presentation_static.py storage_s3.py elasticsearch_api.py sql_console.py hdfs_api.py ssh_terminal.py . +COPY main.py lab_context.py agent_terminal.py workload.py node_registry.py node_ops.py topology_views.py supervisor.py approval_service.py db.py dockhand_envs.py presentation.py database_inventory.py presentation_upload.py presentation_static.py storage_s3.py elasticsearch_api.py sql_console.py hdfs_api.py ssh_terminal.py pipeline_ops.py . RUN mkdir -p /data ENV DATABASE_URL=sqlite:////data/atc-agents.db EXPOSE 3201 diff --git a/api/main.py b/api/main.py index 15140a9..1f2d9a0 100644 --- a/api/main.py +++ b/api/main.py @@ -38,6 +38,7 @@ from presentation_upload import ( from presentation_static import get_static_deck, list_static_decks from storage_s3 import router as storage_s3_router from hdfs_api import router as hdfs_router +from pipeline_ops import router as pipeline_router from elasticsearch_api import router as elasticsearch_router from sql_console import router as sql_router from ssh_terminal import ssh_session @@ -719,6 +720,7 @@ async def lifespan(app: FastAPI): app = FastAPI(title="ATC Command Center API", lifespan=lifespan) app.include_router(storage_s3_router) app.include_router(hdfs_router) +app.include_router(pipeline_router) app.include_router(elasticsearch_router) app.include_router(sql_router) app.add_middleware( diff --git a/api/pipeline_ops.py b/api/pipeline_ops.py new file mode 100644 index 0000000..d953966 --- /dev/null +++ b/api/pipeline_ops.py @@ -0,0 +1,165 @@ +"""Pipeline / data-generation control for the Command Center. + +Triggers per-database Airflow DAGs (light, configurable row counts), reports +run status, and returns live "in sync" counts via Trino so the UI can show the +whole data flow being pulsed and recognised downstream. +""" + +from __future__ import annotations + +import os +import time +from typing import Any + +import httpx +from fastapi import APIRouter, Body, Query +from fastapi.responses import JSONResponse + +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", "") +TRINO_URL = os.getenv("TRINO_URL", "http://10.0.21.50:8089").rstrip("/") +TRINO_USER = os.getenv("TRINO_USER", "mo") + +router = APIRouter(prefix="/api/pipeline", tags=["pipeline"]) + +# UI source key -> Airflow DAG id +SOURCE_DAG = { + "postgres": "gen_postgres", + "mysql": "gen_mysql", + "mongodb": "gen_mongodb", + "cassandra": "gen_cassandra", + "neo4j": "gen_neo4j", + "all": "generate_data_all_databases", + "hadoop": "gen_hadoop_history", +} + +# UI source key -> Trino fully-qualified table for live row counts +SOURCE_COUNT_SQL = { + "postgres": "SELECT count(*) FROM postgres_sales.public.sales_orders", + "mysql": "SELECT count(*) FROM mysql_hr.hr.employee_events", + "mongodb": "SELECT count(*) FROM mongodb_supplychain.supplychain.events", + "cassandra": "SELECT count(*) FROM cassandra_telemetry.telemetry.device_metrics", +} + +_token_cache: dict[str, Any] = {"token": None, "exp": 0.0} + + +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 # tokens last ~24h; refresh well before + return tok + + +async def _trino_scalar(sql: str, timeout: float = 12.0) -> int | None: + try: + async with httpx.AsyncClient(timeout=timeout) as client: + r = await client.post( + f"{TRINO_URL}/v1/statement", + content=sql.encode(), + headers={"X-Trino-User": TRINO_USER}, + ) + data = r.json() + rows: list[Any] = [] + nxt = data.get("nextUri") + if data.get("data"): + rows += data["data"] + while nxt: + rr = await client.get(nxt) + d = rr.json() + if d.get("data"): + rows += d["data"] + if d.get("error"): + return None + nxt = d.get("nextUri") + if rows: + return int(rows[0][0]) + except Exception: + return None + return None + + +@router.post("/generate/{source}") +async def generate(source: str, body: dict[str, Any] = Body(default={})) -> JSONResponse: + dag_id = SOURCE_DAG.get(source) + if not dag_id: + return JSONResponse({"ok": False, "error": f"Unknown source '{source}'"}, status_code=400) + rows = body.get("rows") + conf: dict[str, Any] = {} + if rows is not None: + try: + conf["rows"] = max(1, min(int(rows), 2_000_000)) + except (TypeError, ValueError): + return JSONResponse({"ok": False, "error": "rows must be an integer"}, status_code=400) + try: + async with httpx.AsyncClient() as client: + tok = await _airflow_token(client) + r = await client.post( + f"{AIRFLOW_URL}/api/v2/dags/{dag_id}/dagRuns", + headers={"Authorization": f"Bearer {tok}"}, + json={"logical_date": None, "conf": conf}, + timeout=15, + ) + if r.status_code >= 400: + return JSONResponse({"ok": False, "error": f"Airflow {r.status_code}: {r.text[:300]}"}, status_code=200) + j = r.json() + return JSONResponse({ + "ok": True, + "source": source, + "dag_id": dag_id, + "run_id": j.get("dag_run_id"), + "state": j.get("state"), + "rows": conf.get("rows"), + }) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=200) + + +@router.get("/runs/{source}") +async def runs(source: str, limit: int = Query(5)) -> JSONResponse: + dag_id = SOURCE_DAG.get(source) + if not dag_id: + return JSONResponse({"ok": False, "error": f"Unknown source '{source}'"}, status_code=400) + try: + async with httpx.AsyncClient() as client: + tok = await _airflow_token(client) + headers = {"Authorization": f"Bearer {tok}"} + r = await client.get( + f"{AIRFLOW_URL}/api/v2/dags/{dag_id}/dagRuns", + headers=headers, + params={"order_by": "-run_after", "limit": limit}, + timeout=12, + ) + runs_list = r.json().get("dag_runs", []) + out = [] + for run in runs_list: + out.append({ + "run_id": run.get("dag_run_id"), + "state": run.get("state"), + "start": run.get("start_date"), + "end": run.get("end_date"), + "conf": run.get("conf"), + }) + return JSONResponse({"ok": True, "source": source, "dag_id": dag_id, "runs": out}) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=200) + + +@router.get("/sync") +async def sync(source: str | None = Query(None)) -> JSONResponse: + """Live row counts per source (via Trino) for the 'in sync' display.""" + targets = [source] if source and source in SOURCE_COUNT_SQL else list(SOURCE_COUNT_SQL.keys()) + counts: dict[str, Any] = {} + for src in targets: + counts[src] = await _trino_scalar(SOURCE_COUNT_SQL[src]) + return JSONResponse({"ok": True, "counts": counts}) diff --git a/docker-compose.yml b/docker-compose.yml index e1cc2f7..840b13d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -39,6 +39,8 @@ services: LLM_API_KEY: sk-local LAKEHOUSE_HOST: 10.0.21.50 AIRFLOW_URL: http://10.0.21.55:8080 + TRINO_URL: http://10.0.21.50:8089 + TRINO_USER: ${TRINO_USER:-mo} KAFKA_UI_URL: http://10.0.21.36:9000 HDFS_NN_URL: http://10.0.21.61:9870 S3_ENDPOINT: ${S3_ENDPOINT:-http://10.0.20.111:9020}