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

我想了解在 FastAPI 应用程序中处理 websocket 端点异常的推荐方法是什么。

我尝试过：

app.add_exception_handler(Exception, handle_generic_exception)

它捕获 Exception ，但它没有捕获，例如 ValueError | ||我也尝试过使用

但它似乎不适用于 websocket。 @app.middleware("http") 你们中的任何人都知道处理 websocket 端点的

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")

通用 异常的正确方法是什么吗？ |编辑 2024-05-28： 我终于添加了装饰器来处理我的 websocket 端点的异常。但如果有任何更优雅的解决方案，请告诉我！

EDIT 2024-05-28:

I finally added the decorator for handling exceptions for my websocket endpoints. But if there is any more elegant solution, please let me know!

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)

你对于使用中间件处理 FastAPI WebSocket 端点中的异常的理解是正确的，它们不适用于此场景。中间件主要用于处理 HTTP 请求和响应。

在 WebSocket 端点中，你需要在函数内部处理异常。以下是一些你可以采取的方法：

1. 使用 try-except 块：

这是处理 WebSocket 端点异常最直接的方法。你可以在你的 WebSocket 函数内部使用 try-except 块来捕获和处理异常。

@app.websocket("/ws")
async def ws(websocket: WebSocket):
    await websocket.accept()
    try:
        # 处理 WebSocket 消息
        while True:
            data = await websocket.receive_text()
            # 处理数据并可能引发异常
            if data == "raise":
                raise ValueError("Invalid data received")
            await websocket.send_text(f"Message received: {data}")
    except ValueError as e:
        await websocket.send_text(f"Error: {e}")
    except Exception as e:
        await websocket.send_text(f"An unexpected error occurred: {e}")
    finally:
        await websocket.close()

2. 使用装饰器来处理异常：

你可以创建一个装饰器来处理 WebSocket 函数中的异常，从而避免在每个函数中都编写 try-except 块。

from functools import wraps

def websocket_exception_handler(func):
    @wraps(func)
    async def wrapper(websocket: WebSocket):
        await websocket.accept()
        try:
            await func(websocket)
        except Exception as e:
            await websocket.send_text(f"An error occurred: {e}")
        finally:
            await websocket.close()
    return wrapper

@app.websocket("/ws")
@websocket_exception_handler
async def ws(websocket: WebSocket):
    # 处理 WebSocket 消息
    while True:
        data = await websocket.receive_text()
        if data == "raise":
            raise ValueError("Invalid data received")
        await websocket.send_text(f"Message received: {data}")

3. 自定义异常处理类:

你可以定义自定义异常类，并在 try-except 块中使用它们来处理特定类型的异常。

class WebSocketError(Exception):
    def __init__(self, message, status_code):
        super().__init__(message)
        self.status_code = status_code

@app.websocket("/ws")
async def ws(websocket: WebSocket):
    await websocket.accept()
    try:
        # 处理 WebSocket 消息并可能引发异常
    except ValueError as e:
        await websocket.send_text(f"Error: {e}")
        await websocket.close(code=1002)  # 使用合适的关闭代码
    except WebSocketError as e:
        await websocket.send_text(f"Error: {e}")
        await websocket.close(code=e.status_code)

使用哪种方法取决于你的应用程序的具体需求。

请记住，在处理 WebSocket 端点中的异常时，向客户端发送有意义的错误消息非常重要，以便他们了解发生了什么。 你还可以使用不同的 WebSocket 关闭代码来表示错误的原因.