Added websocket endpoints for filament and vendors

spoolman/api/v1/filament.py

@@ -1,9 +1,10 @@
 """Filament related endpoints."""
 
+import asyncio
 import logging
 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
@@ -16,6 +17,7 @@ from spoolman.database import filament
 from spoolman.database.database import get_db_session
 from spoolman.database.utils import SortOrder
 from spoolman.exceptions import ItemDeleteError
+from spoolman.ws import websocket_manager
 
 logger = logging.getLogger(__name__)
 
@@ -232,7 +234,10 @@ async def find(
 @router.get(
     "/{filament_id}",
     name="Get filament",
-    description="Get a specific filament.",
+    description=(
+        "Get a specific filament. A websocket is served on the same path to listen for changes to the filament. "
+        "The response model is the same for the websocket messages as for this endpoint."
+    ),
     response_model_exclude_none=True,
     responses={404: {"model": Message}},
 )
@@ -244,6 +249,25 @@ async def get(
     return Filament.from_db(db_item)
 
 
+@router.websocket(
+    "/{filament_id}",
+    name="Listen to filament changes",
+)
+async def notify(
+    websocket: WebSocket,
+    filament_id: int,
+) -> None:
+    await websocket.accept()
+    websocket_manager.connect(("filament", str(filament_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(("filament", str(filament_id)), websocket)
+
+
 @router.post(
     "",
     name="Add filament",

spoolman/api/v1/vendor.py

@@ -1,8 +1,9 @@
 """Vendor related endpoints."""
 
+import asyncio
 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
@@ -14,6 +15,7 @@ from spoolman.api.v1.models import Message, Vendor
 from spoolman.database import vendor
 from spoolman.database.database import get_db_session
 from spoolman.database.utils import SortOrder
+from spoolman.ws import websocket_manager
 
 router = APIRouter(
     prefix="/vendor",
@@ -100,7 +102,10 @@ async def find(
 @router.get(
     "/{vendor_id}",
     name="Get vendor",
-    description="Get a specific vendor.",
+    description=(
+        "Get a specific vendor. A websocket is served on the same path to listen for changes to the vendor. "
+        "The response model is the same for the websocket messages as for this endpoint."
+    ),
     response_model_exclude_none=True,
     responses={404: {"model": Message}},
 )
@@ -112,6 +117,25 @@ async def get(
     return Vendor.from_db(db_item)
 
 
+@router.websocket(
+    "/{vendor_id}",
+    name="Listen to vendor changes",
+)
+async def notify(
+    websocket: WebSocket,
+    vendor_id: int,
+) -> None:
+    await websocket.accept()
+    websocket_manager.connect(("vendor", str(vendor_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(("vendor", str(vendor_id)), websocket)
+
+
 @router.post(
     "",
     name="Add vendor",

spoolman/database/filament.py

@@ -10,6 +10,7 @@ from sqlalchemy.exc import IntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import contains_eager, joinedload
 
+from spoolman.api.v1.models import Filament
 from spoolman.database import models, vendor
 from spoolman.database.utils import (
     SortOrder,
@@ -19,6 +20,7 @@ from spoolman.database.utils import (
     parse_nested_field,
 )
 from spoolman.exceptions import ItemDeleteError, ItemNotFoundError
+from spoolman.ws import websocket_manager
 
 
 async def create(
@@ -148,6 +150,7 @@ async def update(
         else:
             setattr(filament, k, v)
     await db.commit()
+    await filament_changed(filament)
     return filament
 
 
@@ -183,3 +186,8 @@ async def find_article_numbers(
     stmt = select(models.Filament.article_number).distinct()
     rows = await db.execute(stmt)
     return [row[0] for row in rows.all() if row[0] is not None]
+
+
+async def filament_changed(filament: models.Filament) -> None:
+    """Notify websocket clients that a filament has changed."""
+    await websocket_manager.send(("filament", str(filament.id)), Filament.from_db(filament).json())

spoolman/database/vendor.py

@@ -6,9 +6,11 @@ from typing import Optional
 from sqlalchemy import func, select
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from spoolman.api.v1.models import Vendor
 from spoolman.database import models
 from spoolman.database.utils import SortOrder, add_where_clause_str
 from spoolman.exceptions import ItemNotFoundError
+from spoolman.ws import websocket_manager
 
 
 async def create(
@@ -87,6 +89,7 @@ async def update(
     for k, v in data.items():
         setattr(vendor, k, v)
     await db.commit()
+    await vendor_changed(vendor)
     return vendor
 
 
@@ -94,3 +97,8 @@ async def delete(db: AsyncSession, vendor_id: int) -> None:
     """Delete a vendor object."""
     vendor = await get_by_id(db, vendor_id)
     await db.delete(vendor)
+
+
+async def vendor_changed(vendor: models.Vendor) -> None:
+    """Notify websocket clients that a vendor has changed."""
+    await websocket_manager.send(("vendor", str(vendor.id)), Vendor.from_db(vendor).json())