wip

learn_sql_model/api/websocket.py

@@ -1,10 +1,11 @@
 from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
 from fastapi.responses import HTMLResponse
 from sqlmodel import Session
+from websockets.exceptions import ConnectionClosed
 
 from learn_sql_model.api.websocket_connection_manager import manager
 from learn_sql_model.config import get_session
-from learn_sql_model.models.hero import Heros
+from learn_sql_model.models.hero import HeroUpdate, Heros
 
 web_socket_router = APIRouter()
 
@@ -81,6 +82,31 @@ async def websocket_endpoint(
     session: Session = Depends(get_session),
 ):
     await websocket.accept()
-    while True:
-        heros = Heros.list(session=session)
-        await websocket.send_text(heros.json())
+    try:
+        while True:
+            heros = Heros.list(session=session)
+            await websocket.send_text(heros.json())
+    except WebSocketDisconnect:
+        print("disconnected")
+    except ConnectionClosed:
+        print("connection closed")
+
+
+@web_socket_router.websocket("/ws-hero-update")
+async def websocket_endpoint(
+    websocket: WebSocket,
+    session: Session = Depends(get_session),
+):
+    await websocket.accept()
+    try:
+        while True:
+            data = await websocket.receive_text()
+            hero = HeroUpdate.parse_raw(data)
+            print(hero)
+            hero.update(session=session)
+            print("hero is updated")
+
+    except WebSocketDisconnect:
+        print("disconnected")
+    except ConnectionClosed:
+        print("connection closed")