Added WebSocket notify API for spools

spoolman/api/v1/spool.py

@@ -1,10 +1,11 @@
 """Spool related endpoints."""
 
+import asyncio
 import logging
 from datetime import datetime
 from typing import Annotated, Optional
 
-from fastapi import APIRouter, Depends, Query
+from fastapi import APIRouter, Depends, Query, WebSocket, WebSocketDisconnect
 from fastapi.encoders import jsonable_encoder
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
@@ -17,6 +18,7 @@ from spoolman.database import spool
 from spoolman.database.database import get_db_session
 from spoolman.database.utils import SortOrder
 from spoolman.exceptions import ItemCreateError
+from spoolman.ws import websocket_manager
 
 logger = logging.getLogger(__name__)
 
@@ -242,7 +244,10 @@ async def find(
 @router.get(
     "/{spool_id}",
     name="Get spool",
-    description="Get a specific spool.",
+    description=(
+        "Get a specific spool. A websocket is served on the same path to listen for changes to the spool. "
+        "The response model is the same for the websocket messages as for this endpoint."
+    ),
     response_model_exclude_none=True,
     responses={404: {"model": Message}},
 )
@@ -254,6 +259,25 @@ async def get(
     return Spool.from_db(db_item)
 
 
+@router.websocket(
+    "/{spool_id}",
+    name="Listen to spool changes",
+)
+async def notify(
+    websocket: WebSocket,
+    spool_id: int,
+) -> None:
+    await websocket.accept()
+    websocket_manager.connect(("spool", str(spool_id)), websocket)
+    try:
+        while True:
+            await asyncio.sleep(0.5)
+            if await websocket.receive_text():
+                await websocket.send_json({"status": "healthy"})
+    except WebSocketDisconnect:
+        websocket_manager.disconnect(("spool", str(spool_id)), websocket)
+
+
 @router.post(
     "",
     name="Add spool",

spoolman/database/spool.py

@@ -10,6 +10,7 @@ from sqlalchemy.exc import NoResultFound
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import contains_eager, joinedload
 
+from spoolman.api.v1.models import Spool
 from spoolman.database import filament, models
 from spoolman.database.utils import (
     SortOrder,
@@ -21,6 +22,7 @@ from spoolman.database.utils import (
 )
 from spoolman.exceptions import ItemCreateError, ItemNotFoundError
 from spoolman.math import weight_from_length
+from spoolman.ws import websocket_manager
 
 
 def utc_timezone_naive(dt: datetime) -> datetime:
@@ -179,6 +181,7 @@ async def update(
         else:
             setattr(spool, k, v)
     await db.commit()
+    await spool_changed(spool)
     return spool
 
 
@@ -231,6 +234,7 @@ async def use_weight(db: AsyncSession, spool_id: int, weight: float) -> models.S
     spool.last_used = datetime.utcnow()
 
     await db.commit()
+    await spool_changed(spool)
     return spool
 
 
@@ -275,6 +279,7 @@ async def use_length(db: AsyncSession, spool_id: int, length: float) -> models.S
     spool.last_used = datetime.utcnow()
 
     await db.commit()
+    await spool_changed(spool)
     return spool
 
 
@@ -296,3 +301,8 @@ async def find_lot_numbers(
     stmt = sqlalchemy.select(models.Spool.lot_nr).distinct()
     rows = await db.execute(stmt)
     return [row[0] for row in rows.all() if row[0] is not None]
+
+
+async def spool_changed(spool: models.Spool) -> None:
+    """Notify websocket clients that a spool has changed."""
+    await websocket_manager.send(("spool", str(spool.id)), Spool.from_db(spool).json())

spoolman/ws.py

@@ -0,0 +1,48 @@
+"""Websocket functionality."""
+
+import logging
+
+from fastapi import WebSocket
+
+logger = logging.getLogger(__name__)
+
+
+class WebsocketManager:
+    """Websocket manager."""
+
+    def __init__(self) -> None:
+        """Initialize."""
+        self.connections: dict[str, set[WebSocket]] = {}
+
+    def connect(self, pool: tuple[str, ...], websocket: WebSocket) -> None:
+        """Connect a websocket."""
+        pool_str = ",".join(pool)
+        if pool_str not in self.connections:
+            self.connections[pool_str] = set()
+        self.connections[pool_str].add(websocket)
+        logger.info(
+            "Client %s is now listening on pool %s",
+            websocket.client.host if websocket.client else "?",
+            pool_str,
+        )
+
+    def disconnect(self, pool: tuple[str, ...], websocket: WebSocket) -> None:
+        """Disconnect a websocket."""
+        pool_str = ",".join(pool)
+        if pool_str in self.connections:
+            self.connections[pool_str].remove(websocket)
+        logger.info(
+            "Client %s has stopped listening on pool %s",
+            websocket.client.host if websocket.client else "?",
+            pool_str,
+        )
+
+    async def send(self, pool: tuple[str, ...], message: str) -> None:
+        """Send a message to all websockets in a pool."""
+        pool_str = ",".join(pool)
+        if pool_str in self.connections:
+            for websocket in self.connections[pool_str]:
+                await websocket.send_text(message)
+
+
+websocket_manager = WebsocketManager()