Add SSE updates

This commit is contained in:
Andras Schmelczer 2026-06-04 22:12:10 +01:00
parent d1732128e2
commit 688bc0cfe9
13 changed files with 958 additions and 42 deletions

View file

@ -1,12 +1,16 @@
from __future__ import annotations
import asyncio
import json
import sqlite3
import time
from typing import Annotated
import structlog
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi import APIRouter, Depends, Header, HTTPException, Request
from fastapi.responses import StreamingResponse
from . import events
from .auth import get_current_user
from .db import db_connection
from .limits import limiter
@ -18,6 +22,7 @@ from .models import (
HealthResponse,
HslColor,
PageOut,
PutDataResponse,
RegisterRequest,
RegisterResponse,
TowerOut,
@ -26,6 +31,34 @@ from .models import (
router = APIRouter(prefix="/api/v1")
logger = structlog.get_logger(__name__)
# How long the SSE loop waits for a new revision before emitting a keepalive
# comment. Keeps the connection warm through proxies and bounds disconnect
# detection latency.
SSE_KEEPALIVE_SECONDS = 20
def _read_revision(conn: sqlite3.Connection, user_id: str) -> int:
row = conn.execute(
"SELECT revision FROM users WHERE id = ?", (user_id,)
).fetchone()
return int(row["revision"]) if row is not None else 0
def _parse_if_match(value: str | None) -> int | None:
"""Parse the client's CAS base revision from the If-Match header.
Tolerates an ETag-style quoted form (``"5"``). Returns None when the header
is absent or unparseable, in which case the write proceeds unguarded
(last-writer-wins) this keeps older cached clients working across a deploy.
"""
if value is None:
return None
cleaned = value.strip().strip('"')
try:
return int(cleaned)
except ValueError:
return None
@router.get("/health", response_model=HealthResponse)
async def health() -> HealthResponse:
@ -63,6 +96,8 @@ async def get_data(
user_id: Annotated[str, Depends(get_current_user)],
) -> DataOut:
with db_connection() as conn:
revision = _read_revision(conn, user_id)
pages_rows = conn.execute(
"""
SELECT id, name, hide_create_tower_button, keep_tasks_open, default_date_from, default_date_to
@ -140,24 +175,44 @@ async def get_data(
)
)
return DataOut(pages=pages_out)
return DataOut(pages=pages_out, revision=revision)
@router.put("/data", status_code=204)
@router.put("/data", response_model=PutDataResponse)
@limiter.limit("30/minute")
async def put_data(
request: Request,
body: DataIn,
user_id: Annotated[str, Depends(get_current_user)],
) -> None:
if_match: Annotated[str | None, Header(alias="If-Match")] = None,
) -> PutDataResponse:
# Tree-replace semantics mean the request size IS the user's total storage,
# so the 2 MiB request cap (enforced by payload_size_middleware) is the only
# quota we need.
now = int(time.time())
base_revision = _parse_if_match(if_match)
with db_connection() as conn:
conn.execute("BEGIN IMMEDIATE")
try:
current_revision = _read_revision(conn, user_id)
# Compare-and-swap: reject if another client advanced the revision
# under us. The client refetches, merges, and retries with the
# fresh base. An absent If-Match opts out of the guard.
if base_revision is not None and base_revision != current_revision:
conn.rollback()
raise HTTPException(
status_code=409,
detail={
"error": "conflict",
"detail": "Revision is stale; refetch and retry",
"revision": current_revision,
},
)
new_revision = current_revision + 1
# Delete existing data for this user (cascades to towers + blocks)
conn.execute("DELETE FROM pages WHERE user_id = ?", (user_id,))
@ -231,13 +286,16 @@ async def put_data(
),
)
# Update last_seen_at
# Bump the revision and refresh last_seen_at in the same transaction.
conn.execute(
"UPDATE users SET last_seen_at = ? WHERE id = ?",
(now, user_id),
"UPDATE users SET last_seen_at = ?, revision = ? WHERE id = ?",
(now, new_revision, user_id),
)
conn.commit()
except HTTPException:
# Already rolled back (e.g. the CAS 409); just propagate.
raise
except sqlite3.IntegrityError as exc:
conn.rollback()
raise HTTPException(
@ -248,4 +306,66 @@ async def put_data(
conn.rollback()
raise
logger.info("data_replaced", user_id=token_log_id(user_id), pages=len(body.pages))
# Notify other live clients for this user to refetch. Done after commit so
# subscribers never observe a revision that isn't durably stored yet.
events.publish(user_id, new_revision)
logger.info(
"data_replaced",
user_id=token_log_id(user_id),
pages=len(body.pages),
revision=new_revision,
)
return PutDataResponse(revision=new_revision)
def _format_revision_event(revision: int) -> str:
return f"event: revision\ndata: {json.dumps({'revision': revision})}\n\n"
@router.get("/events")
async def events_stream(
request: Request,
user_id: Annotated[str, Depends(get_current_user)],
) -> StreamingResponse:
"""Server-Sent Events stream that notifies a client to refetch.
Authenticated with the same Bearer token as the rest of the API clients
consume it via fetch()+ReadableStream (EventSource cannot set headers), so
the token never leaks into a URL. Each event carries the current revision;
the client refetches GET /data when it sees a revision newer than its own.
"""
queue = events.subscribe(user_id)
async def event_generator():
try:
# Emit the current revision immediately so a (re)connecting client
# can catch up on anything it missed while disconnected.
with db_connection() as conn:
yield _format_revision_event(_read_revision(conn, user_id))
while True:
try:
revision = await asyncio.wait_for(
queue.get(), timeout=SSE_KEEPALIVE_SECONDS
)
yield _format_revision_event(revision)
except asyncio.TimeoutError:
if await request.is_disconnected():
break
# Comment line: keeps proxies from idling the connection out.
yield ": keepalive\n\n"
finally:
events.unsubscribe(user_id, queue)
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
# Disable nginx response buffering for this stream even if the
# location block forgets `proxy_buffering off`.
"X-Accel-Buffering": "no",
},
)

View file

@ -0,0 +1,77 @@
"""In-process pub/sub for Server-Sent Events.
Single-worker deployment (see Dockerfile: uvicorn with no ``--workers``), so a
plain in-memory registry is enough there is one event loop and every
subscriber queue lives on it. If this ever grows to multiple workers, this
module is the seam to swap for a cross-process bus (Redis pub/sub, LISTEN/NOTIFY,
or polling the ``users.revision`` column).
Each subscriber gets a 1-slot queue that always holds the *latest* revision to
notify. Coalescing is intentional: the SSE payload is just "something changed,
current revision is N", so a burst of writes collapses to a single refetch
signal rather than a backlog.
"""
from __future__ import annotations
import asyncio
import structlog
logger = structlog.get_logger(__name__)
# user_id -> set of per-connection queues. Each queue carries the most recent
# revision the connection still needs to flush to its client.
_subscribers: dict[str, set[asyncio.Queue[int]]] = {}
def subscribe(user_id: str) -> asyncio.Queue[int]:
"""Register a new SSE connection for ``user_id`` and return its queue.
Must be called from the event loop (the SSE route handler is async), so the
queue binds to the running loop.
"""
queue: asyncio.Queue[int] = asyncio.Queue(maxsize=1)
_subscribers.setdefault(user_id, set()).add(queue)
return queue
def unsubscribe(user_id: str, queue: asyncio.Queue[int]) -> None:
"""Drop a connection's queue; prune the user entry when the last one goes."""
subs = _subscribers.get(user_id)
if subs is None:
return
subs.discard(queue)
if not subs:
_subscribers.pop(user_id, None)
def publish(user_id: str, revision: int) -> None:
"""Notify every live connection for ``user_id`` of the new revision.
Called from the (single) event loop right after a PUT commits, so the
``put_nowait`` calls are loop-safe. Coalesces into each 1-slot queue: if a
connection has not yet drained its previous signal, replace it with the
newer revision rather than block or grow unbounded.
"""
subs = _subscribers.get(user_id)
if not subs:
return
for queue in subs:
try:
queue.put_nowait(revision)
except asyncio.QueueFull:
# Connection is behind — drop the stale value and keep the newest.
try:
queue.get_nowait()
except asyncio.QueueEmpty:
pass
try:
queue.put_nowait(revision)
except asyncio.QueueFull:
pass
def connection_count(user_id: str) -> int:
"""Number of live SSE connections for a user (used in tests)."""
return len(_subscribers.get(user_id, ()))

View file

@ -0,0 +1,7 @@
-- Per-user revision counter for multi-client sync.
-- Bumped inside the same transaction as every PUT /data so clients can:
-- * detect another client's write (cheap compare, no full-tree diff), and
-- * guard their own writes with compare-and-swap (stale If-Match -> 409).
-- Existing users start at 0; their next confirmed GET seeds the client base.
ALTER TABLE users ADD COLUMN revision INTEGER NOT NULL DEFAULT 0;

View file

@ -127,6 +127,14 @@ class DataIn(BaseModel):
class DataOut(BaseModel):
pages: list[PageOut]
# Monotonic per-user version, bumped on every successful PUT. Clients keep
# it as their compare-and-swap base and to detect another client's writes.
revision: int
class PutDataResponse(BaseModel):
# The new revision after this write; the client adopts it as its CAS base.
revision: int
class RegisterRequest(BaseModel):