api: pipeline_ops router (per-source generate trigger, run status, Trino sync counts)
This commit is contained in:
mo
+1
-1
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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})
|
||||
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user