From 92c91c2594cb7d4890e3fc4596003a704175788b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=86=D0=9C-24=20=D0=92=D0=BB=D0=B0=D0=B4=D0=B8=D1=81?=
 =?UTF-8?q?=D0=BB=D0=B0=D0=B2=20=D0=9A=D0=BE=D0=B2=D0=B0=D0=BB=D0=B5=D0=BD?=
 =?UTF-8?q?=D0=BA=D0=BE?= <kovalenkov222@gmail.com>
Date: Thu, 26 Mar 2026 13:38:04 +0000
Subject: [PATCH] refactor: use try block for error handling and session
 rollback

---
 store/main.py | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/store/main.py b/store/main.py
index 7030c4e..0a72d5e 100644
--- a/store/main.py
+++ b/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)