diff --git a/store/main.py b/store/main.py index 6639e41..5069286 100644 --- a/store/main.py +++ b/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) diff --git a/store/schemas.py b/store/schemas.py index e224570..1771d18 100644 --- a/store/schemas.py +++ b/store/schemas.py @@ -50,6 +50,6 @@ class ProcessedAgentData(BaseModel): road_state: str agent_data: AgentData -class WebSockerData(BaseModel): +class WebSocketData(BaseModel): id: int