spb/company-atlas

1"""Server-sent events for the live feed: poll `events` every `POLL_S` with a `(detected_at, id)` cursor, heartbeat every `HEARTBEAT_S`.23The first message is always a heartbeat carrying the cursor so clients (and proxies) see bytes immediately; reconnecting clients4pass the last cursor back as `?since=`. `max_s` bounds the connection so upstream proxies with idle timeouts can reconnect cleanly.5"""6from __future__ import annotations78import asyncio9import logging10from collections.abc import AsyncIterator11from datetime import datetime12from typing import Any1314import orjson15from fastapi import Request16from sse_starlette.sse import EventSourceResponse1718from companyatlas.api import queries as q19from companyatlas.api import serializers as ser20from companyatlas.db import connection2122log = logging.getLogger("companyatlas.api.sse")2324POLL_S = 5.025HEARTBEAT_S = 20.026BATCH = 5027MAX_STREAM_S = 360028SSE_HEADERS = {"cache-control": "no-store, no-transform", "x-accel-buffering": "no", "connection": "keep-alive"}293031def _dump(payload: Any) -> str:32    return orjson.dumps(payload, option=orjson.OPT_UTC_Z, default=str).decode()333435async def poll_new_events(since: datetime | None, last_id: str | None, *, event_type: str | None = None, min_importance: float | None = None,36                          limit: int = BATCH) -> list[dict[str, Any]]:37    where, params = q.event_filters(event_type=event_type, min_importance=min_importance)38    if since is not None:39        if last_id:40            where.append("(e.detected_at, e.id) > (:cur_at, :cur_id)")41            params.update(cur_at=since, cur_id=last_id)42        else:43            where.append("e.detected_at > :cur_at")44            params["cur_at"] = since45    async with connection() as conn:46        rows = await q.fetch_events(conn, where, params, sort="recent", limit=limit)47    return [ser.event(r) for r in reversed(rows)]      # oldest first so clients append in order484950async def live_event_stream(request: Request | None, *, since: datetime | None, event_type: str | None = None,51                            min_importance: float | None = None, max_s: float = MAX_STREAM_S) -> AsyncIterator[dict[str, Any]]:52    loop = asyncio.get_running_loop()53    started = loop.time()54    cursor_at: datetime | None = since or q.now_utc()55    cursor_id: str | None = None56    last_sent = loop.time()57    yield {"event": "heartbeat", "data": _dump({"at": q.now_utc(), "cursor": cursor_at})}58    while True:59        if request is not None and await request.is_disconnected():60            return61        try:62            events = await poll_new_events(cursor_at, cursor_id, event_type=event_type, min_importance=min_importance)63        except Exception:64            log.warning("live stream poll failed", exc_info=True)65            events = []66        for ev in events:67            cursor_at, cursor_id = ev["detected_at"], ev["id"]68            last_sent = loop.time()69            yield {"event": "event", "id": ev["id"], "data": _dump(ev)}70        now = loop.time()71        if now - last_sent >= HEARTBEAT_S:72            last_sent = now73            yield {"event": "heartbeat", "data": _dump({"at": q.now_utc(), "cursor": cursor_at})}74        remaining = max_s - (now - started)75        if remaining <= 0:76            yield {"event": "end", "data": _dump({"cursor": cursor_at, "reason": "max_s reached"})}77            return78        await asyncio.sleep(min(POLL_S, remaining))798081def sse_response(generator: AsyncIterator[dict[str, Any]]) -> EventSourceResponse:82    return EventSourceResponse(generator, headers=SSE_HEADERS, ping=HEARTBEAT_S * 3)838485__all__ = ["HEARTBEAT_S", "MAX_STREAM_S", "POLL_S", "live_event_stream", "poll_new_events", "sse_response"]86