159 lines
6.1 KiB
Python
159 lines
6.1 KiB
Python
"""WebSocket API endpoints"""
|
|
|
|
import asyncio
|
|
import logging
|
|
from datetime import UTC, datetime
|
|
from typing import Optional
|
|
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
|
|
from jose import jwt, JWTError
|
|
|
|
from app.core.config import settings
|
|
from app.core.time import to_iso8601_utc
|
|
from app.core.websocket.manager import manager
|
|
from app.core.websocket.ue_scene import expand_scene_payload_for_transport, ue_scene_state_store
|
|
from app.db.session import async_session_factory
|
|
|
|
logger = logging.getLogger(__name__)
|
|
router = APIRouter()
|
|
SUPPORTED_CHANNELS = [
|
|
"gpu_clusters",
|
|
"submarine_cables",
|
|
"ixp_nodes",
|
|
"alerts",
|
|
"dashboard",
|
|
"datasource_tasks",
|
|
"ue_scene",
|
|
]
|
|
|
|
|
|
async def authenticate_token(token: str) -> Optional[dict]:
|
|
"""Authenticate WebSocket connection via token"""
|
|
try:
|
|
payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
|
|
if payload.get("type") != "access":
|
|
logger.warning(f"WebSocket auth failed: wrong token type")
|
|
return None
|
|
return payload
|
|
except JWTError as e:
|
|
logger.warning(f"WebSocket auth failed: {e}")
|
|
return None
|
|
|
|
|
|
@router.websocket("/ws")
|
|
async def websocket_endpoint(
|
|
websocket: WebSocket,
|
|
token: str = Query(...),
|
|
):
|
|
"""WebSocket endpoint for real-time data"""
|
|
logger.info(f"WebSocket connection attempt with token: {token[:20]}...")
|
|
payload = await authenticate_token(token)
|
|
if payload is None:
|
|
logger.warning("WebSocket authentication failed, closing connection")
|
|
await websocket.close(code=4001)
|
|
return
|
|
|
|
user_id = str(payload.get("sub"))
|
|
await manager.connect(websocket, user_id)
|
|
|
|
try:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "connection_established",
|
|
"timestamp": to_iso8601_utc(datetime.now(UTC)),
|
|
"data": {
|
|
"connection_id": f"conn_{user_id}",
|
|
"server_version": settings.VERSION,
|
|
"heartbeat_interval": 30,
|
|
"supported_channels": SUPPORTED_CHANNELS,
|
|
},
|
|
}
|
|
)
|
|
|
|
while True:
|
|
try:
|
|
data = await asyncio.wait_for(websocket.receive_json(), timeout=30)
|
|
|
|
if data.get("type") == "heartbeat":
|
|
await websocket.send_json(
|
|
{
|
|
"type": "heartbeat",
|
|
"timestamp": to_iso8601_utc(datetime.now(UTC)),
|
|
"data": {"action": "pong", "timestamp": to_iso8601_utc(datetime.now(UTC))},
|
|
}
|
|
)
|
|
elif data.get("type") == "subscribe":
|
|
requested_channels = data.get("data", {}).get("channels", [])
|
|
channels = manager.subscribe(websocket, requested_channels)
|
|
await websocket.send_json(
|
|
{
|
|
"type": "subscription_confirmed",
|
|
"timestamp": to_iso8601_utc(datetime.now(UTC)),
|
|
"data": {"action": "subscribe", "channels": channels},
|
|
}
|
|
)
|
|
elif data.get("type") == "sync_request":
|
|
sync_data = data.get("data", {})
|
|
channel = sync_data.get("channel")
|
|
if channel == "ue_scene":
|
|
async with async_session_factory() as session:
|
|
scene_payloads = await ue_scene_state_store.get_sync_payloads(
|
|
session,
|
|
last_sequence=sync_data.get("last_sequence"),
|
|
reason=sync_data.get("reason"),
|
|
)
|
|
for scene_payload in scene_payloads:
|
|
for transport_payload in expand_scene_payload_for_transport(scene_payload):
|
|
await websocket.send_json(
|
|
{
|
|
"type": "data_frame",
|
|
"channel": "ue_scene",
|
|
"timestamp": to_iso8601_utc(datetime.now(UTC)),
|
|
"data": transport_payload,
|
|
}
|
|
)
|
|
else:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "error",
|
|
"timestamp": to_iso8601_utc(datetime.now(UTC)),
|
|
"data": {
|
|
"message": f"Unsupported sync channel: {channel}",
|
|
},
|
|
}
|
|
)
|
|
elif data.get("type") == "control_frame":
|
|
await websocket.send_json(
|
|
{
|
|
"type": "control_acknowledged",
|
|
"timestamp": to_iso8601_utc(datetime.now(UTC)),
|
|
"data": {
|
|
"target": data.get("data", {}).get("target"),
|
|
"command": data.get("data", {}).get("command"),
|
|
"accepted": True,
|
|
},
|
|
}
|
|
)
|
|
else:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "ack",
|
|
"timestamp": to_iso8601_utc(datetime.now(UTC)),
|
|
"data": {"received": True},
|
|
}
|
|
)
|
|
|
|
except asyncio.TimeoutError:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "heartbeat",
|
|
"timestamp": to_iso8601_utc(datetime.now(UTC)),
|
|
"data": {"action": "ping"},
|
|
}
|
|
)
|
|
|
|
except WebSocketDisconnect:
|
|
pass
|
|
finally:
|
|
manager.disconnect(websocket, user_id)
|