store/main.py

@@ -60,6 +60,7 @@ async def websocket_endpoint(websocket: WebSocket, user_id: int):
         # receive forever
         while True:
             data = await websocket.receive_text()
+            try:
                 if (data):
                     ws_data = WebSocketData.model_validate(json.loads(data))
                     session = SessionLocal()
@@ -70,9 +71,12 @@ async def websocket_endpoint(websocket: WebSocket, user_id: int):
                     ).returning(processed_agent_data)
                     res = session.execute(update_query).fetchone()
                     if (not res):
-                    print("Websocket update fail")
+                        session.rollback()
+                        raise Exception("Error while websocket PUT")
                     session.commit()
+            finally:
                 session.close()                  
+         
     except WebSocketDisconnect:
         subscriptions.remove(websocket)