store/main.py

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