diff --git a/store/docker/db/structure.sql b/store/docker/db/structure.sql index 06b6f76..01566a3 100644 --- a/store/docker/db/structure.sql +++ b/store/docker/db/structure.sql @@ -7,5 +7,6 @@ CREATE TABLE processed_agent_data ( z FLOAT, latitude FLOAT, longitude FLOAT, - timestamp TIMESTAMP -); \ No newline at end of file + timestamp TIMESTAMP, + visible BOOLEAN +); diff --git a/store/main.py b/store/main.py index 94ea9f7..8ebef39 100644 --- a/store/main.py +++ b/store/main.py @@ -8,12 +8,13 @@ from sqlalchemy import ( Integer, String, Float, + Boolean, DateTime, ) 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() @@ -30,6 +31,7 @@ processed_agent_data = Table( Column("latitude", Float), Column("longitude", Float), Column("timestamp", DateTime), + Column("visible", Boolean), ) # WebSocket subscriptions @@ -57,7 +59,24 @@ async def websocket_endpoint(websocket: WebSocket, user_id: int): # receive forever while True: - await websocket.receive_text() + data = await websocket.receive_text() + 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) @@ -81,6 +100,7 @@ def ProcessedAgentData_to_td(data: List[ProcessedAgentData]): "latitude": item.agent_data.gps.latitude, "longitude": item.agent_data.gps.longitude, "timestamp": item.agent_data.timestamp, + "visible": True, } for item in data ] diff --git a/store/schemas.py b/store/schemas.py index 3d13dee..1771d18 100644 --- a/store/schemas.py +++ b/store/schemas.py @@ -49,3 +49,7 @@ class AgentData(BaseModel): class ProcessedAgentData(BaseModel): road_state: str agent_data: AgentData + +class WebSocketData(BaseModel): + id: int +