scenarios when using FastAPI with WebSockets and SQLAlchemy for real-time notifications
I'm refactoring my project and I tried several approaches but none seem to work..... I'm trying to implement a real-time notification system using FastAPI with WebSockets and SQLAlchemy. However, I'm working with an scenario where the WebSocket connection closes unexpectedly after I send a message. I'm using FastAPI version 0.68.0 and SQLAlchemy 1.4.27. Here’s a simplified version of my code: ```python from fastapi import FastAPI, WebSocket from sqlalchemy import create_engine, Column, Integer, String from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker DATABASE_URL = "sqlite:///./test.db" engine = create_engine(DATABASE_URL) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base = declarative_base() class Notification(Base): __tablename__ = "notifications" id = Column(Integer, primary_key=True, index=True) message = Column(String, index=True) Base.metadata.create_all(bind=engine) app = FastAPI() @app.websocket("/ws/notifications") async def websocket_endpoint(websocket: WebSocket): await websocket.accept() db = SessionLocal() while True: data = await websocket.receive_text() new_notification = Notification(message=data) db.add(new_notification) db.commit() await websocket.send_text(f"Notification sent: {data}") # Close the connection unexpectedly if data == "close": break await websocket.close() db.close() ``` When I send a message through the WebSocket, it works as expected, but if I send 'close', the connection closes, but I get this behavior in my logs: ``` behavior: Exception in ASGI application Traceback (most recent call last): File ".../site-packages/starlette/routing.py", line 259, in __call__ await self.app(scope, receive, send) File ".../site-packages/starlette/middleware/errors.py", line 159, in __call__ raise exc File ".../site-packages/starlette/routing.py", line 257, in __call__ await self.app(scope, receive, send) File ".../site-packages/fastapi/applications.py", line 128, in __call__ await super().__call__(scope, receive, send) File ".../site-packages/starlette/applications.py", line 106, in __call__ await self.middleware_stack(scope, receive, send) File ".../site-packages/starlette/middleware/middleware.py", line 22, in __call__ await self.dispatch(scope, receive, send) File ".../site-packages/starlette/middleware/errors.py", line 159, in __call__ raise exc File ".../site-packages/starlette/routing.py", line 257, in __call__ await self.app(scope, receive, send) File ".../site-packages/starlette/routing.py", line 258, in __call__ await self.app(scope, receive, send) File ".../site-packages/starlette/routing.py", line 257, in __call__ await self.app(scope, receive, send) File ".../site-packages/fastapi/applications.py", line 128, in __call__ await super().__call__(scope, receive, send) File ".../site-packages/starlette/applications.py", line 106, in __call__ await self.middleware_stack(scope, receive, send) File ".../site-packages/starlette/middleware/middleware.py", line 22, in __call__ await self.dispatch(scope, receive, send) File ".../site-packages/starlette/middleware/errors.py", line 159, in __call__ raise exc File ".../site-packages/starlette/routing.py", line 257, in __call__ await self.app(scope, receive, send) File ".../site-packages/starlette/routing.py", line 258, in __call__ await self.app(scope, receive, send) ``` I’m unsure if this is a question with how I’m handling the database session or if it's a WebSocket scenario. I’ve tried adding behavior handling and ensuring the database session is closed properly, but the behavior still occurs. Any suggestions on how to resolve this scenario or improve the implementation for stability would be appreciated! I'm working on a web app that needs to handle this. For context: I'm using Python on Linux. What's the best practice here? Could this be a known issue? I appreciate any insights!