Add initial Websocket

The websocket is used to send immediate notification for newly created sensor data / state

Author: mawaeg

api/main.py

@@ -1,19 +1,23 @@
+import asyncio
 from contextlib import asynccontextmanager
 from typing import Annotated
 
 from fastapi import Depends, FastAPI
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from api.routers import authentication, forecast, permissions, sensors, serverstats, users
+from api.routers import authentication, forecast, permissions, sensors, serverstats, users, websocket
 from api.utils.database import dispose_database, get_session
 from api.utils.security import get_current_user
+from api.utils.websocket_connection_handler import get_websocket_handler
 
 
 @asynccontextmanager
 async def lifespan(_app: FastAPI):
+    asyncio.get_event_loop().create_task(get_websocket_handler().event_loop())
     yield
     await dispose_database()
 
+
 app: FastAPI = FastAPI(root_path="/weatherapi", lifespan=lifespan)
 
 app.include_router(authentication.auth_router)
@@ -22,6 +26,7 @@ async def lifespan(_app: FastAPI):
 app.include_router(forecast.forecast_router)
 app.include_router(permissions.permissions_router)
 app.include_router(serverstats.serverstats_router)
+app.include_router(websocket.websocket_router)
 
 
 @app.get("/")

api/routers/sensors.py

@@ -1,6 +1,6 @@
 from typing import Annotated
 
-from fastapi import APIRouter, Depends, status
+from fastapi import APIRouter, BackgroundTasks, Depends, status
 from sqlalchemy.ext.asyncio.session import AsyncSession
 from sqlalchemy.sql import func
 from sqlmodel import NUMERIC, cast, or_, select
@@ -22,6 +22,7 @@
 from api.utils.permissions import get_user_read_permissions, get_user_write_permissions
 from api.utils.security import get_current_superuser, get_current_user
 from api.utils.sensor_utils import get_is_valid_sensor_type, get_sensor_from_db
+from api.utils.websocket_connection_handler import WebsocketHandler, get_websocket_handler
 
 sensors_router = APIRouter(tags=["Sensors"], prefix="/sensor")
 
@@ -117,6 +118,8 @@ async def create_sensor(
 @sensors_router.post("/{sensor_id}/data", response_model=SensorData, status_code=status.HTTP_201_CREATED)
 async def create_sensor_data(
     session: Annotated[AsyncSession, Depends(get_session)],
+    ws_handler: Annotated[WebsocketHandler, Depends(get_websocket_handler)],
+    background_tasks: BackgroundTasks,
     current_user: Annotated[DBUser, Depends(get_current_user)],
     sensor_id: int,
     data: SensorDataCreate,
@@ -140,6 +143,9 @@ async def create_sensor_data(
     session.add(data)
     await session.commit()
     await session.refresh(data)
+
+    background_tasks.add_task(ws_handler.add_event, data)
+
     return data
 
 
@@ -210,6 +216,8 @@ async def get_sensor_data_daily(
 @sensors_router.post("/{sensor_id}/state", response_model=SensorState, status_code=status.HTTP_201_CREATED)
 async def create_sensor_state(
     session: Annotated[AsyncSession, Depends(get_session)],
+    ws_handler: Annotated[WebsocketHandler, Depends(get_websocket_handler)],
+    background_tasks: BackgroundTasks,
     current_user: Annotated[DBUser, Depends(get_current_user)],
     sensor_id: int,
     data: SensorStateCreate,
@@ -233,6 +241,9 @@ async def create_sensor_state(
     session.add(data)
     await session.commit()
     await session.refresh(data)
+
+    background_tasks.add_task(ws_handler.add_event, data)
+
     return data
 
 

api/routers/websocket.py

@@ -0,0 +1,25 @@
+from typing import Annotated
+
+from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
+
+from api.models.database_models import DBUser
+from api.utils.security import get_current_user_ws
+from api.utils.websocket_connection_handler import WebsocketHandler, get_websocket_handler
+
+websocket_router = APIRouter()
+
+
+@websocket_router.websocket("/ws")
+async def websocket_endpoint(
+    websocket: WebSocket,
+    current_user: Annotated[DBUser, Depends(get_current_user_ws)],
+    ws_handler: Annotated[WebsocketHandler, Depends(get_websocket_handler)],
+):
+    await websocket.accept()
+    ws_handler.add(current_user, websocket)
+    print(f"Accepted connection with {current_user.username}")
+    try:
+        while True:
+            await websocket.receive_text()
+    except WebSocketDisconnect:
+        ws_handler.remove(current_user)

api/utils/security.py

@@ -1,7 +1,7 @@
 from datetime import datetime, timedelta, timezone
 from typing import Annotated
 
-from fastapi import Depends
+from fastapi import Depends, Header, WebSocketException, status
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 from jose import JWTError, jwt
 from passlib.context import CryptContext
@@ -92,6 +92,7 @@ def create_access_token(data: dict, expires_delta: timedelta | None = None):
 async def get_current_user(
     token: Annotated[str, Depends(oauth2_scheme)], session: Annotated[AsyncSession, Depends(get_session)]
 ) -> DBUser:
+    # ToDo Handle expired token somehow
     try:
         payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
         username: str = payload.get("sub")
@@ -105,6 +106,29 @@ async def get_current_user(
     return user
 
 
+async def get_current_user_ws(
+    session: Annotated[AsyncSession, Depends(get_session)], authorization: Annotated[str | None, Header()] = None
+) -> DBUser:
+    if authorization is None:
+        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
+
+    if not authorization.startswith("Bearer "):
+        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
+
+    token = authorization.split(" ")[1]
+    try:
+        payload = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
+        username: str = payload.get("sub")
+        if username is None:
+            raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
+    except JWTError:
+        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
+    user = await get_user(username, session)
+    if user is None:
+        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
+    return user
+
+
 async def get_current_superuser(user: Annotated[DBUser, Depends(get_current_user)]):
     if not user.superuser:
         raise MISSING_PRIVILEGES

api/utils/websocket_connection_handler.py

@@ -0,0 +1,49 @@
+import asyncio
+
+from fastapi import HTTPException, WebSocket
+from pydantic import BaseModel
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+from api.models.database_models import DatabaseModelBase, DBUser
+from api.utils.database import get_engine, get_session
+from api.utils.permissions import get_user_read_permissions
+
+
+class WebsocketHandler:
+    def __init__(self):
+        self._connections: dict[int, tuple[DBUser, WebSocket]] = {}
+        self._message_queue: asyncio.Queue[DatabaseModelBase] = asyncio.Queue()
+
+    def add(self, user: DBUser, websocket: WebSocket) -> bool:
+        if self._connections.get(user.id):
+            return False
+
+        self._connections[user.id] = (user, websocket)
+
+    def remove(self, user: DBUser) -> bool:
+        if self._connections.pop(user.id, None):
+            return True
+        return False
+
+    async def event_loop(self):
+        engine: AsyncEngine = await get_engine()
+        while True:
+            event = await self._message_queue.get()
+            async for session in get_session(engine):
+                for user, websocket in self._connections.values():
+                    try:
+                        await get_user_read_permissions(session, user, event.id)
+                        await websocket.send_json(event.model_dump_json())
+                    except HTTPException:
+                        pass
+                self._message_queue.task_done()
+
+    async def add_event(self, data: BaseModel):
+        await self._message_queue.put(data)
+
+
+_websocket_handler = WebsocketHandler()
+
+
+def get_websocket_handler() -> WebsocketHandler:
+    return _websocket_handler

pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     "sqlalchemy==2.0.25",
     "sqlmodel==0.0.14",
     "uvicorn==0.27.0",
+    "websockets>=15.0.1",
 ]
 
 [tool.black]