From 3e2bbed17eb42ca4918effafc95257ead83f4888 Mon Sep 17 00:00:00 2001 From: Donkie Date: Sun, 15 Oct 2023 20:36:48 +0200 Subject: [PATCH] Added more websocket endpoints Can now subscribe to: - spool/filament/vendor level events, not specific to any specific ID - root level events, captures all events sent --- spoolman/api/v1/filament.py | 30 ++++++++++++++++++++++++++++-- spoolman/api/v1/router.py | 31 ++++++++++++++++++++++++++++++- spoolman/api/v1/spool.py | 30 ++++++++++++++++++++++++++++-- spoolman/api/v1/vendor.py | 30 ++++++++++++++++++++++++++++-- 4 files changed, 114 insertions(+), 7 deletions(-) diff --git a/spoolman/api/v1/filament.py b/spoolman/api/v1/filament.py index 86e76db..fc85fe0 100644 --- a/spoolman/api/v1/filament.py +++ b/spoolman/api/v1/filament.py @@ -110,9 +110,17 @@ class FilamentUpdateParameters(FilamentParameters): @router.get( "", name="Find filaments", - description="Get a list of filaments that matches the search query.", + description=( + "Get a list of filaments that matches the search query. " + "A websocket is served on the same path to listen for updates to any filament, or added or deleted filaments. " + "See the HTTP Response code 299 for the content of the websocket messages." + ), response_model_exclude_none=True, - response_model=list[Filament], + responses={ + 200: {"model": list[Filament]}, + 404: {"model": Message}, + 299: {"model": FilamentEvent, "description": "Websocket message"}, + }, ) async def find( *, @@ -231,6 +239,24 @@ async def find( ) +@router.websocket( + "", + name="Listen to filament changes", +) +async def notify_any( + websocket: WebSocket, +) -> None: + await websocket.accept() + websocket_manager.connect(("filament",), 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",), websocket) + + @router.get( "/{filament_id}", name="Get filament", diff --git a/spoolman/api/v1/router.py b/spoolman/api/v1/router.py index 57ca61c..04fbc13 100644 --- a/spoolman/api/v1/router.py +++ b/spoolman/api/v1/router.py @@ -2,9 +2,10 @@ # ruff: noqa: D103 +import asyncio import logging -from fastapi import FastAPI +from fastapi import FastAPI, WebSocket, WebSocketDisconnect from fastapi.responses import JSONResponse from starlette.requests import Request from starlette.responses import Response @@ -12,6 +13,7 @@ from starlette.responses import Response from spoolman import env from spoolman.database.database import backup_global_db from spoolman.exceptions import ItemNotFoundError +from spoolman.ws import websocket_manager from . import filament, models, other, spool, vendor @@ -20,6 +22,15 @@ logger = logging.getLogger(__name__) app = FastAPI( title="Spoolman REST API v1", version="1.0.0", + description=""" + REST API for Spoolman. + + The API is served on the path `/api/v1/`. + + Some endpoints also serve a websocket on the same path. The websocket is used to listen for changes to the data + that the endpoint serves. The websocket messages are JSON objects. Additionally, there is a root-level websocket + endpoint that listens for changes to any data in the database. + """, ) @@ -73,6 +84,24 @@ async def backup(): # noqa: ANN201 return models.BackupResponse(path=str(path)) +@app.websocket( + "/", + name="Listen to any changes", +) +async def notify( + websocket: WebSocket, +) -> None: + await websocket.accept() + websocket_manager.connect((), 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((), websocket) + + # Add routers app.include_router(filament.router) app.include_router(spool.router) diff --git a/spoolman/api/v1/spool.py b/spoolman/api/v1/spool.py index 86e6bb3..f9d083c 100644 --- a/spoolman/api/v1/spool.py +++ b/spoolman/api/v1/spool.py @@ -68,9 +68,17 @@ class SpoolUseParameters(BaseModel): @router.get( "", name="Find spool", - description="Get a list of spools that matches the search query.", + description=( + "Get a list of spools that matches the search query. " + "A websocket is served on the same path to listen for updates to any spool, or added or deleted spools. " + "See the HTTP Response code 299 for the content of the websocket messages." + ), response_model_exclude_none=True, - response_model=list[Spool], + responses={ + 200: {"model": list[Spool]}, + 404: {"model": Message}, + 299: {"model": SpoolEvent, "description": "Websocket message"}, + }, ) async def find( *, @@ -245,6 +253,24 @@ async def find( ) +@router.websocket( + "", + name="Listen to spool changes", +) +async def notify_any( + websocket: WebSocket, +) -> None: + await websocket.accept() + websocket_manager.connect(("spool",), 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",), websocket) + + @router.get( "/{spool_id}", name="Get spool", diff --git a/spoolman/api/v1/vendor.py b/spoolman/api/v1/vendor.py index 4d70849..e3b72fe 100644 --- a/spoolman/api/v1/vendor.py +++ b/spoolman/api/v1/vendor.py @@ -46,9 +46,17 @@ class VendorUpdateParameters(VendorParameters): @router.get( "", name="Find vendor", - description="Get a list of vendors that matches the search query.", + description=( + "Get a list of vendors that matches the search query. " + "A websocket is served on the same path to listen for updates to any vendor, or added or deleted vendors. " + "See the HTTP Response code 299 for the content of the websocket messages." + ), response_model_exclude_none=True, - response_model=list[Vendor], + responses={ + 200: {"model": list[Vendor]}, + 404: {"model": Message}, + 299: {"model": VendorEvent, "description": "Websocket message"}, + }, ) async def find( db: Annotated[AsyncSession, Depends(get_db_session)], @@ -99,6 +107,24 @@ async def find( ) +@router.websocket( + "", + name="Listen to vendor changes", +) +async def notify_any( + websocket: WebSocket, +) -> None: + await websocket.accept() + websocket_manager.connect(("vendor",), 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",), websocket) + + @router.get( "/{vendor_id}", name="Get vendor",