Compare commits

..

3 commits

Author SHA1 Message Date
Johan Lundberg
c3cb93d903 Merge branch 'master' into lundberg_data_generator 2025-05-02 02:19:39 +02:00
erik
4f9fdb911e Websocket test 2025-05-01 19:26:00 +00:00
13dac398c5 Merge pull request 'add make command make reformat' (#1) from lundberg_add_black into master
Reviewed-on: #1
2025-05-01 19:19:25 +00:00

40
main.py
View file

@ -3,7 +3,7 @@ import json
import sqlite3
from typing import Dict
from fastapi import FastAPI, Header, HTTPException, Query
from fastapi import FastAPI, Header, HTTPException, Query, WebSocket, WebSocketDisconnect
from fastapi.responses import JSONResponse
from fastapi.routing import APIRoute
from fastapi.staticfiles import StaticFiles
@ -11,6 +11,8 @@ from pydantic import BaseModel
from typing import Optional
from db import init_db, save_snapshot, DB_FILE
import asyncio
from starlette.concurrency import run_in_threadpool
# ------------------------------------------------------------------
app = FastAPI()
@ -194,6 +196,42 @@ def get_trails(
]
return JSONResponse(content={"trails": trails})
# -------------------- WebSocket endpoints -----------------------
browser_conns: set[WebSocket] = set()
async def _broadcast_to_browser_clients(snapshot: dict):
for ws in list(browser_conns):
try:
await ws.send_json(snapshot)
except WebSocketDisconnect:
browser_conns.remove(ws)
@app.websocket("/ws/position")
async def ws_receive_snapshots(websocket: WebSocket, secret: str = Query(...)):
await websocket.accept()
if secret != SHARED_SECRET:
await websocket.close(code=1008)
return
try:
while True:
data = await websocket.receive_json()
snap = TelemetrySnapshot.parse_obj(data)
live_snapshots[snap.character_name] = snap.dict()
await run_in_threadpool(save_snapshot, snap.dict())
await _broadcast_to_browser_clients(snap.dict())
except WebSocketDisconnect:
pass
@app.websocket("/ws/live")
async def ws_live_updates(websocket: WebSocket):
await websocket.accept()
browser_conns.add(websocket)
try:
while True:
await asyncio.sleep(3600)
except WebSocketDisconnect:
browser_conns.remove(websocket)
# -------------------- static frontend ---------------------------
app.mount("/", StaticFiles(directory="static", html=True), name="static")