diff --git a/api/pipeline_ops.py b/api/pipeline_ops.py index d953966..37812da 100644 --- a/api/pipeline_ops.py +++ b/api/pipeline_ops.py @@ -7,6 +7,7 @@ whole data flow being pulsed and recognised downstream. from __future__ import annotations +import asyncio import os import time from typing import Any @@ -43,6 +44,29 @@ SOURCE_COUNT_SQL = { } _token_cache: dict[str, Any] = {"token": None, "exp": 0.0} +# count(*) over these connectors does a full scan (slow), so a background task +# refreshes the counts periodically and the endpoint always returns instantly. +_sync_cache: dict[str, Any] = {"counts": {k: None for k in SOURCE_COUNT_SQL}, "ts": 0.0} +_refresher_started = False +_SYNC_INTERVAL = 60.0 + + +async def _sync_refresh_loop() -> None: + while True: + try: + keys = list(SOURCE_COUNT_SQL.keys()) + results = await asyncio.gather( + *[_trino_scalar(SOURCE_COUNT_SQL[k], deadline_s=90.0) for k in keys] + ) + merged = dict(_sync_cache["counts"]) + for k, v in zip(keys, results): + if v is not None: + merged[k] = v + _sync_cache["counts"] = merged + _sync_cache["ts"] = time.time() + except Exception: + pass + await asyncio.sleep(_SYNC_INTERVAL) async def _airflow_token(client: httpx.AsyncClient) -> str: @@ -61,9 +85,16 @@ async def _airflow_token(client: httpx.AsyncClient) -> str: return tok -async def _trino_scalar(sql: str, timeout: float = 12.0) -> int | None: +async def _trino_scalar(sql: str, deadline_s: float = 8.0) -> int | None: + """Run a scalar Trino query with a hard wall-clock deadline. + + count(*) over a large Cassandra table can scan for a long time; without a + total deadline the result-paging loop would hang the endpoint. On timeout + we abort the query and return None so the UI degrades gracefully. + """ + end = time.time() + deadline_s try: - async with httpx.AsyncClient(timeout=timeout) as client: + async with httpx.AsyncClient(timeout=5.0) as client: r = await client.post( f"{TRINO_URL}/v1/statement", content=sql.encode(), @@ -71,10 +102,16 @@ async def _trino_scalar(sql: str, timeout: float = 12.0) -> int | None: ) data = r.json() rows: list[Any] = [] - nxt = data.get("nextUri") if data.get("data"): rows += data["data"] + nxt = data.get("nextUri") while nxt: + if time.time() > end: + try: + await client.delete(nxt, timeout=3.0) + except Exception: + pass + return None rr = await client.get(nxt) d = rr.json() if d.get("data"): @@ -156,10 +193,19 @@ async def runs(source: str, limit: int = Query(5)) -> JSONResponse: @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}) +async def sync(source: str | None = Query(None), refresh: bool = Query(False)) -> JSONResponse: + """Live row counts per source (via Trino) for the 'in sync' display. + + Counts are cached (TTL) and refreshed concurrently because count(*) over + these connectors does a full scan. Returns immediately from cache when + fresh; otherwise refreshes once (other concurrent callers reuse the cache). + """ + global _refresher_started + if not _refresher_started: + _refresher_started = True + asyncio.create_task(_sync_refresh_loop()) + counts = _sync_cache["counts"] + age = round(time.time() - _sync_cache["ts"], 1) if _sync_cache["ts"] else None + if source and source in counts: + return JSONResponse({"ok": True, "counts": {source: counts[source]}, "age_s": age}) + return JSONResponse({"ok": True, "counts": counts, "age_s": age})