"""FastAPI app — POST /events writes to primary, GET /ws/events reads from

the secondary-fed broadcast hub.

Lifecycle:

Startup → open ONE primary connection + ONE secondary connection

(the latter inside the Poller). Both connections live for

the lifetime of the process. Mutual-RA-TLS handshakes

happen here, NOT per request.

Shutdown → close primary connection, stop poller (which closes the

secondary connection). Forwarders die cleanly.

Reconnect: if the primary connection raises OperationalError on a write,

we close it and reopen synchronously inside the request handler. The

request returns 503 — the client retries, and the next request lands

on the freshly-handshaken connection. Same recover-on-error contract

as the poller.

Routes:

GET / minimal HTML UI (form + WebSocket + log)

GET /events recent rows, newest-first, query ?limit=N (default 100, max 1000)

POST /events body {topic, payload} → INSERT → returns {id}

GET /ws/events WebSocket → newline-delimited JSON event stream

GET /healthz simple liveness + connection state

"""

from __future__ import annotations

import asyncio

import json

import logging

from collections.abc import AsyncIterator

from contextlib import asynccontextmanager

from pathlib import Path

from typing import Any

import psycopg

from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect

from fastapi.responses import FileResponse, JSONResponse

from pydantic import BaseModel, Field

from .broadcast import Hub

from .config import Settings, load_settings

from .connection import open_primary

from .poller import Poller

log = logging.getLogger(__name__)

def _find_static_dir() -> Path:

"""Locate the static dir whether running from source tree or installed.

In the source tree, `static/` is at the repo root (three levels above

this file). In the Docker image, the Dockerfile sets `WORKDIR /app`

and copies `static/` to `/app/static`, but the package gets installed

into site-packages — so the source-tree heuristic doesn't apply. Try

both layouts and fall back to cwd."""

candidates = [

Path(__file__).resolve().parent.parent.parent / "static",

Path.cwd() / "static",

Path("/app/static"),

]

for c in candidates:

if (c / "index.html").is_file():

return c

return candidates[0]

_STATIC_DIR = _find_static_dir()

class EventIn(BaseModel):

"""POST /events request body. Topic + free-form JSON payload."""

topic: str = Field(..., min_length=1, max_length=128)

payload: dict[str, Any] = Field(default_factory=dict)

class EventOut(BaseModel):

id: int

@asynccontextmanager

async def _lifespan(app: FastAPI) -> AsyncIterator[None]:

"""Open long-lived connections + start the poller. Tear down on exit."""

settings: Settings = app.state.settings

hub: Hub = Hub()

app.state.hub = hub

loop = asyncio.get_running_loop()

log.info("opening primary connection")

primary_conn = await loop.run_in_executor(None, open_primary, settings)

app.state.primary_conn = primary_conn

poller = Poller(settings, hub)

await poller.start()

app.state.poller = poller

log.info(

"ready: cluster=%s primary_role=%s secondary_role=%s poll=%dms",

settings.cluster_uuid,

settings.primary_role,

settings.secondary_role,

settings.poll_interval_ms,

)

try:

yield

finally:

log.info("shutting down")

await poller.stop()

try:

primary_conn.close()

except Exception as e: # noqa: BLE001

log.warning("error closing primary: %s", e)

def create_app(settings: Settings | None = None) -> FastAPI:

"""Construct the FastAPI app. Importable from tests."""

settings = settings or load_settings()

app = FastAPI(

title="teesql-example",

description=(

"Worked example of read/write split + WebSocket fan-out over a "

"teesql cluster. Writes go to the primary; reads stream from the "

"secondary."

),

version="0.1.0",

lifespan=_lifespan,

)

app.state.settings = settings

_register_routes(app)

return app

def _reopen_primary(app: FastAPI) -> psycopg.Connection:

"""Close the dead primary connection, reopen, replace on app.state."""

settings: Settings = app.state.settings

old: psycopg.Connection = app.state.primary_conn

log.warning("reopening primary connection")

try:

old.close()

except Exception: # noqa: BLE001

pass

new = open_primary(settings)

app.state.primary_conn = new

return new

def _register_routes(app: FastAPI) -> None:

@app.get("/", include_in_schema=False)

async def index() -> FileResponse:

return FileResponse(_STATIC_DIR / "index.html")

@app.get("/healthz")

async def healthz(request: Request) -> JSONResponse:

s: Settings = request.app.state.settings

primary_ok = False

try:

conn: psycopg.Connection = request.app.state.primary_conn

with conn.cursor() as cur:

cur.execute("SELECT 1")

cur.fetchone()

conn.commit()

primary_ok = True

except Exception as e: # noqa: BLE001

log.warning("healthz: primary check failed: %s", e)

hub: Hub = request.app.state.hub

return JSONResponse(

{

"ok": primary_ok,

"cluster_uuid": s.cluster_uuid,

"primary_role": s.primary_role,

"secondary_role": s.secondary_role,

"subscribers": hub.subscriber_count,

"poll_interval_ms": s.poll_interval_ms,

}

)

@app.get("/events")

async def get_events(request: Request, limit: int = 100) -> JSONResponse:

"""Return the most recent events, newest-first.

Reads from the long-lived primary connection. We could route this

through the secondary, but the secondary is owned by the poller

which is busy ticking every 50ms; a one-shot history fetch on

the already-open primary is simpler and avoids cross-thread

contention. The point of the read/write split demo is the

STREAMING path (poller → hub → WebSocket), not this one-shot.

"""

if limit < 1 or limit > 1000:

raise HTTPException(status_code=422, detail="limit must be in [1, 1000]")

loop = asyncio.get_running_loop()

def _do_select(conn: psycopg.Connection) -> list[dict[str, Any]]:

with conn.cursor() as cur:

cur.execute(

"SELECT id, topic, payload, recorded_at FROM events "

"ORDER BY id DESC LIMIT %s",

(limit,),

)

rows = list(cur.fetchall())

conn.commit()

return [

{

"id": int(r["id"]),

"topic": r["topic"],

"payload": r["payload"],

"recorded_at": r["recorded_at"].isoformat(),

}

for r in rows

]

try:

data = await loop.run_in_executor(

None, _do_select, request.app.state.primary_conn

)

except psycopg.OperationalError as e:

log.warning("GET /events failed (%s); reopening primary and retrying once", e)

try:

new_conn = await loop.run_in_executor(None, _reopen_primary, request.app)

data = await loop.run_in_executor(None, _do_select, new_conn)

except Exception as e2: # noqa: BLE001

raise HTTPException(status_code=503, detail=f"primary unreachable: {e2}") from e2

return JSONResponse(data)

@app.post("/events", response_model=EventOut)

async def post_event(body: EventIn, request: Request) -> EventOut:

loop = asyncio.get_running_loop()

def _do_insert(conn: psycopg.Connection) -> int:

with conn.cursor() as cur:

cur.execute(

"INSERT INTO events (topic, payload) "

"VALUES (%s, %s::jsonb) RETURNING id",

(body.topic, json.dumps(body.payload)),

)

row = cur.fetchone()

if row is None:

raise RuntimeError("INSERT … RETURNING returned no row")

# row_factory=dict_row → row['id']; cast for safety.

new_id = int(row["id"])

conn.commit()

return new_id

# First attempt; on OperationalError, reopen and retry exactly once.

# Every other exception bubbles up as a 500.

try:

new_id = await loop.run_in_executor(

None, _do_insert, request.app.state.primary_conn

)

except psycopg.OperationalError as e:

log.warning("INSERT failed (%s); reopening primary and retrying once", e)

try:

new_conn = await loop.run_in_executor(None, _reopen_primary, request.app)

new_id = await loop.run_in_executor(None, _do_insert, new_conn)

except Exception as e2: # noqa: BLE001

log.error("INSERT retry failed: %s", e2)

raise HTTPException(status_code=503, detail=f"primary unreachable: {e2}") from e2

return EventOut(id=new_id)

@app.websocket("/ws/events")

async def ws_events(ws: WebSocket) -> None:

"""Fan out every replicated event to this socket as JSON."""

await ws.accept()

hub: Hub = ws.app.state.hub

log.info("ws connected")

try:

async with hub.subscribe() as queue:

while True:

event = await queue.get()

# Send as a single JSON line so the browser can

# parse with `JSON.parse(msg.data)` directly.

await ws.send_text(json.dumps(event))

except WebSocketDisconnect:

log.info("ws disconnected")

except Exception as e: # noqa: BLE001

log.warning("ws closed with error: %s", e)

try:

await ws.close()

except Exception:

pass