IoT-Systems/store/main.py

import json
from typing import Set, Dict, List
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Body
from fastapi.encoders import jsonable_encoder
from sqlalchemy import (
    Table,
    Column,
    Integer,
    String,
    Float,
    DateTime,
)
from sqlalchemy.sql import select

from database import metadata, SessionLocal
from schemas import ProcessedAgentData, ProcessedAgentDataInDB

# FastAPI app setup
app = FastAPI()

processed_agent_data = Table(
    "processed_agent_data",
    metadata,
    Column("id", Integer, primary_key=True, index=True),
    Column("road_state", String),
    Column("user_id", Integer),
    Column("x", Float),
    Column("y", Float),
    Column("z", Float),
    Column("latitude", Float),
    Column("longitude", Float),
    Column("timestamp", DateTime),
)

# WebSocket subscriptions
subscriptions: Dict[int, Set[WebSocket]] = {}


# FastAPI WebSocket endpoint
@app.websocket("/ws/{user_id}")
async def websocket_endpoint(websocket: WebSocket, user_id: int):
    await websocket.accept()

    if user_id not in subscriptions:
        subscriptions[user_id] = set()

    subscriptions[user_id].add(websocket)

    try:
        # send already available data
        r = processed_agent_data.select()
        stored_data = SessionLocal().execute(r).fetchall()

        jsonable_data = [{c.name: getattr(i, c.name) for c in processed_agent_data.columns} for i in stored_data]
        for i in jsonable_data:
            i['timestamp'] = i['timestamp'].strftime("%Y-%m-%dT%H:%M:%SZ")

        await websocket.send_json(json.dumps(jsonable_data))

        # receive forever
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        subscriptions[user_id].remove(websocket)


# Function to send data to subscribed users
async def send_data_to_subscribers(user_id: int, data):
    if user_id in subscriptions:
        for websocket in subscriptions[user_id]:
            await websocket.send_json(json.dumps(data))


# FastAPI CRUDL endpoints


@app.post("/processed_agent_data/")
async def create_processed_agent_data(data: List[ProcessedAgentData], user_id: int = Body(..., embed=True)):
    session = SessionLocal()
    try:
        created_data = [
            {
                "road_state": item.road_state,
                "user_id": user_id,
                "x": item.agent_data.accelerometer.x,
                "y": item.agent_data.accelerometer.y,
                "z": item.agent_data.accelerometer.z,
                "latitude": item.agent_data.gps.latitude,
                "longitude": item.agent_data.gps.longitude,
                "timestamp": item.agent_data.timestamp,
            }
            for item in data
        ]
        stmt = processed_agent_data.insert().values(created_data).returning(processed_agent_data)
        result = session.execute(stmt)
        created_records = [dict(row._mapping) for row in result.fetchall()]
        session.commit()

        for record in created_records:
            await send_data_to_subscribers(user_id, jsonable_encoder(record))
        return created_records
    except Exception as err:
        session.rollback()
        print(f"Database error: {err}")
        raise HTTPException(status_code=500, detail="Internal Server Error")
    finally:
        session.close()


@app.get(
    "/processed_agent_data/{processed_agent_data_id}",
    response_model=ProcessedAgentDataInDB,
)
def read_processed_agent_data(processed_agent_data_id: int):
    session = SessionLocal()
    try:
        stmt = select(processed_agent_data).where(
            processed_agent_data.c.id == processed_agent_data_id
        )
        res = session.execute(stmt).fetchone()
        if not res:
            raise HTTPException(status_code=404, detail="Not found")

        return dict(res._mapping)

    finally:
        session.close()


@app.get("/processed_agent_data/", response_model=list[ProcessedAgentDataInDB])
def list_processed_agent_data():
    session = SessionLocal()
    try:
        stmt = select(processed_agent_data)
        res = session.execute(stmt).fetchall()
        if not res:
            raise HTTPException(status_code=404, detail="Not found")

        return [dict(r._mapping) for r in res]

    finally:
        session.close()


@app.put(
    "/processed_agent_data/{processed_agent_data_id}",
    response_model=ProcessedAgentDataInDB,
)
def update_processed_agent_data(processed_agent_data_id: int, data: ProcessedAgentData):
    # Update data
    session = SessionLocal()

    try:
        query = select(processed_agent_data).where(
            processed_agent_data.c.id == processed_agent_data_id
        )
        result = session.execute(query).fetchone()

        if not result:
            raise HTTPException(status_code=404, detail="Data not found")

        update_query = (
            processed_agent_data.update()
            .where(processed_agent_data.c.id == processed_agent_data_id)
            .values(
                road_state=data.road_state,
                user_id=data.agent_data.user_id,
                x=data.agent_data.accelerometer.x,
                y=data.agent_data.accelerometer.y,
                z=data.agent_data.accelerometer.z,
                latitude=data.agent_data.gps.latitude,
                longitude=data.agent_data.gps.longitude,
                timestamp=data.agent_data.timestamp,
            )
        )

        session.execute(update_query)
        session.commit()

        updated_result = session.execute(query).fetchone()

        return ProcessedAgentDataInDB(**updated_result._mapping)

    finally:
        session.close()


@app.delete(
    "/processed_agent_data/{processed_agent_data_id}",
    response_model=ProcessedAgentDataInDB,
)
def delete_processed_agent_data(processed_agent_data_id: int):
    # Delete by id
    session = SessionLocal()

    try:
        query = select(processed_agent_data).where(
            processed_agent_data.c.id == processed_agent_data_id
        )
        result = session.execute(query).fetchone()

        if not result:
            raise HTTPException(status_code=404, detail="Data not found")

        delete_query = processed_agent_data.delete().where(
            processed_agent_data.c.id == processed_agent_data_id
        )

        session.execute(delete_query)
        session.commit()

        return ProcessedAgentDataInDB(**result._mapping)

    finally:
        session.close()

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=8000)