api: non-blocking /sync via background refresher (full-scan counts cached)

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})