FastAPI - 如何处理 websocket 端点中的通用异常

Question

我想了解在 FastAPI 应用程序中处理

websocket

端点异常的推荐方法是什么。

我尝试过：

app.add_exception_handler(Exception, handle_generic_exception)

它捕获了

Exception

，但它没有捕获，例如

ValueError

。

我也尝试过使用

@app.middleware("http")

但它似乎不适用于 websockets。

from fastapi import FastAPI, Request, Response, WebSocket

app = FastAPI()


app.add_exception_handler(AppException, handle_disconnect)

@app.middleware("http")
async def generic_exception_middleware(request: Request | WebSocket, call_next):
    try:
        return await call_next(request)
    except Exception as exc:
        send_something_to_client()
        print("some exception")


@app.websocket("/ws")
async def ws(websocket: WebSocket):
    raise ValueError("foo")

你们中有人知道处理 websocket 端点的generic异常的正确方法是什么吗？

编辑2024年5月28日：

我最终添加了装饰器来处理我的 websocket 端点的异常。但如果有任何更优雅的解决方案，请告诉我！

def handle_exceptions(func):
    """Decorator for handling exceptions."""

    @functools.wraps(func)
    async def wrapper(websocket: WebSocket):
        try:
            await func(websocket)
        except WebSocketDisconnect as exc:
            raise exc
        except AppError as exc:
            await app_error_handler(websocket, exc)
        except Exception as exc:  # pylint: disable=broad-exception-caught
            await generic_exception_handler(websocket, exc)

    return wrapper

[...]

@handle_exceptions
async def accept_api_v1(websocket: WebSocket):
    """Handle api v1 request."""
    await websocket.accept()
    do_things()

@app.websocket("/api/v1")
async def ws_api(websocket: WebSocket):
    """API v1 endpoint"""
    await accept_api_v1(websocket)

Answer 1

我也在寻找更好的解决方案。目前我将 ASGIMiddleware 添加到 fastapi 应用程序，它捕获 websocket 和 http。如果只想处理 websocket，可以使用scope['type'] == 'websocket' 来过滤请求。

from fastapi import FastAPI, WebSocket


class ASGIMiddleware:
    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        try:
            await self.app(scope, receive, send)
        except:
            do_something()


app = FastAPI()
app.add_middleware(ASGIMiddleware)


@app.websocket('/sample')
async def sample(socket: WebSocket):
    raise ValueError('ERROR')