FastAPI实现数据库新增条目时通过WebSocket推送数据

你的核心问题在于两个方面：一是WebSocket端点的循环逻辑错误导致持续推送，二是SQLAlchemy事件监听器的写法不符合规范，且没有建立事件与WebSocket推送的关联机制。以下是具体的修复步骤：

1. 修正SQLAlchemy事件监听器 #

原事件监听器的参数和逻辑完全错误，after_insert事件的回调函数需接收插入的实例对象，而非数据库会话。我们需要通过一个异步队列将新插入的数据传递给WebSocket管理器。

修改models.py：

from sqlalchemy import Boolean, Column, ForeignKey, Integer, String, Float, DateTime, event
from sqlalchemy.schema import UniqueConstraint
from sqlalchemy.orm import relationship
from .database import Base
import datetime
from pytz import timezone


class Measurement(Base):
    __tablename__ = "measurements"

    id = Column(Integer, primary_key=True, index=True)
    device_key = Column(String(length=40), ForeignKey("devices.device_key"))
    inside_temp = Column(Float)
    outside_temp = Column(Float)
    inside_humidity = Column(Float)
    outside_humidity = Column(Float)
    current_capacity = Column(Float)
    timestamp = Column(DateTime, default=lambda: datetime.datetime.now(tz=timezone('America/Los_Angeles')))

    device = relationship("Device", back_populates="measurements")

    def _as_dict(self):
        return {c.name: getattr(self, c.name) for c in self.__table__.columns}


@event.listens_for(Measurement, "after_insert")
def on_measurement_insert(target, connection, **kwargs):
    # 延迟导入避免循环依赖
    from . import main
    # 将新插入的测量数据放入异步队列
    main.manager.new_data_queue.put_nowait(target)


class Device(Base):
    __tablename__ = "devices"
    device_key = Column(String(length=40), unique=True, primary_key=True)
    name = Column(String)
    hardware = Column(String)
    firmware = Column(String)
    software = Column(String)

    measurements = relationship("Measurement", back_populates="device")

2. 实现WebSocket管理器与推送逻辑 #

添加WebSocket管理器来管理活跃连接，并通过异步队列监听数据库事件，实现仅在数据更新时推送。

修改main.py：

from fastapi import FastAPI, Depends, WebSocket, WebSocketDisconnect
from sql_app.database import engine, Session
from sql_app import models
from fastapi.encoders import jsonable_encoder
import sql_app.schemas as schemas
from sql_app.database import Base, get_db
import datetime
import uvicorn
import asyncio


class WebSocketManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []
        self.new_data_queue = asyncio.Queue()

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, data: dict):
        # 给所有活跃连接推送新数据
        for connection in self.active_connections:
            await connection.send_json(data)

    async def listen_for_updates(self):
        # 持续监听队列，有新数据就广播
        while True:
            new_measurement = await self.new_data_queue.get()
            await self.broadcast(jsonable_encoder(new_measurement))


manager = WebSocketManager()

# 启动时启动监听任务
@app.on_event("startup")
async def startup():
    asyncio.create_task(manager.listen_for_updates())

def create_tables():
    print("Creating Tables..")
    Base.metadata.create_all(bind=engine)

app = FastAPI()


@app.post("/measurement/")
async def create_measurement(measurement: schemas.MeasurementCreate, db: Session = Depends(get_db)):
    new_measurement = models.Measurement(**measurement.dict())
    db.add(new_measurement)
    db.commit()
    db.refresh(new_measurement)
    return new_measurement


@app.post("/create_device/")
async def create_device(device: schemas.DeviceCreate, db: Session = Depends(get_db)):
    new_device = models.Device(
        device_key=device.device_key,
        name=device.name,
        hardware=device.hardware,
        firmware=device.firmware,
        software=device.software
    )
    db.add(new_device)
    db.commit()
    return db.refresh(new_device)


@app.get("/measurement/")
async def get_measurements(db: Session = Depends(get_db)):
    return db.query(models.Measurement).filter(
        models.Measurement.timestamp >= datetime.datetime.now() - datetime.timedelta(days=30)
    ).all()


@app.websocket("/ws")
async def dashboard_data(websocket: WebSocket, db: Session = Depends(get_db)):
    await manager.connect(websocket)
    # 发送初始的30天数据
    initial_data = db.query(models.Measurement).filter(
        models.Measurement.timestamp >= datetime.datetime.now() - datetime.timedelta(days=30)
    ).all()
    await websocket.send_json(jsonable_encoder(initial_data))
    
    try:
        # 保持连接，等待断开事件
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        manager.disconnect(websocket)


if __name__ == "__main__":
    uvicorn.run(app, host="localhost", port=8000)

修复逻辑说明 #

1. 事件监听器修正：after_insert事件触发时，将新插入的Measurement实例放入异步队列，而非直接返回数据。
2. WebSocket管理器：
   • 管理所有活跃的WebSocket连接
   • 启动独立任务监听队列，一旦有新数据就广播给所有客户端
   • 客户端连接时发送初始历史数据，断开时清理连接
3. 避免无效循环：原WebSocket端点的无限循环持续调用measurement_stream导致频繁推送，现在改为被动等待断开事件，由管理器统一处理推送。

内容的提问来源于stack exchange，提问作者Sean Palmer