store/main.py

@@ -14,7 +14,7 @@ from sqlalchemy import (
 from sqlalchemy.sql import select
 
 from database import metadata, SessionLocal
-from schemas import ProcessedAgentData, ProcessedAgentDataInDB
+from schemas import ProcessedAgentData, ProcessedAgentDataInDB, WebSocketData
 
 # FastAPI app setup
 app = FastAPI()
@@ -59,7 +59,20 @@ async def websocket_endpoint(websocket: WebSocket, user_id: int):
 
         # receive forever
         while True:
-            await websocket.receive_text()
+            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()
     except WebSocketDisconnect:
         subscriptions.remove(websocket)
 

store/schemas.py

@@ -50,6 +50,6 @@ class ProcessedAgentData(BaseModel):
     road_state: str
     agent_data: AgentData
     
-class WebSockerData(BaseModel):
+class WebSocketData(BaseModel):
     id: int