From 3c5b3cea5ac19f5084561c25ddd89af33f7fe3e7 Mon Sep 17 00:00:00 2001 From: Donkie Date: Sun, 15 Oct 2023 15:39:39 +0200 Subject: [PATCH] Added websocket endpoints for filament and vendors --- spoolman/api/v1/filament.py | 28 ++++++++++++++++++++++++++-- spoolman/api/v1/vendor.py | 28 ++++++++++++++++++++++++++-- spoolman/database/filament.py | 8 ++++++++ spoolman/database/vendor.py | 8 ++++++++ 4 files changed, 68 insertions(+), 4 deletions(-) diff --git a/spoolman/api/v1/filament.py b/spoolman/api/v1/filament.py index 817817f..ea10d43 100644 --- a/spoolman/api/v1/filament.py +++ b/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", diff --git a/spoolman/api/v1/vendor.py b/spoolman/api/v1/vendor.py index 4a69640..3d6e649 100644 --- a/spoolman/api/v1/vendor.py +++ b/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", diff --git a/spoolman/database/filament.py b/spoolman/database/filament.py index 6d5ebef..9611193 100644 --- a/spoolman/database/filament.py +++ b/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()) diff --git a/spoolman/database/vendor.py b/spoolman/database/vendor.py index 2619f5b..6a56b5f 100644 --- a/spoolman/database/vendor.py +++ b/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())