FastAPIWebSocket依赖注入吗

wen python案例 2026-07-11 262

本文目录导读：

主要依赖注入方式
重要注意事项
完整示例：带认证的聊天室

是的，FastAPI 的 WebSocket 支持依赖注入，但使用方式与普通 HTTP 路由有些区别。

主要依赖注入方式

基本依赖注入（最常用）

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends
from typing import Optional
app = FastAPI()
# 定义一个依赖
async def get_token(
    websocket: WebSocket, 
    token: Optional[str] = None
) -> str:
    # 从查询参数中获取 token
    if token is None:
        # 也可以从 headers 或其他地方获取
        token = websocket.headers.get("authorization", "").replace("Bearer ", "")
    if not token:
        await websocket.close(code=1008)  # 策略违规
        raise WebSocketDisconnect()
    return token
# 使用依赖注入
@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    token: str = Depends(get_token)
):
    await websocket.accept()
    # 使用 token
    print(f"User authenticated with token: {token}")
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Message received: {data}")
    except WebSocketDisconnect:
        print("Client disconnected")

使用 `Annotated` 语法（Python 3.9+）

from typing import Annotated
async def verify_token(websocket: WebSocket) -> str:
    token = websocket.headers.get("authorization", "")
    if not token:
        await websocket.close(code=1008)
        raise WebSocketDisconnect()
    return token
@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    token: Annotated[str, Depends(verify_token)]
):
    await websocket.accept()
    # ... 处理逻辑

类依赖（支持状态管理）

from fastapi import Depends
from typing import Optional
class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []
    async def connect(self, websocket: WebSocket, user_id: str):
        await websocket.accept()
        self.active_connections.append(websocket)
        return {"status": "connected", "user_id": user_id}
    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)
async def get_manager():
    manager = ConnectionManager()
    try:
        yield manager
    finally:
        # 清理逻辑
        pass
@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    manager: ConnectionManager = Depends(get_manager)
):
    user_id = "user123"
    result = await manager.connect(websocket, user_id)
    # ... 处理逻辑

多依赖链式注入

from fastapi import FastAPI, WebSocket, Depends
from typing import Annotated
# 层级依赖
async def verify_token(websocket: WebSocket) -> str:
    token = websocket.headers.get("authorization", "")
    if not token:
        await websocket.close(code=1008)
        raise WebSocketDisconnect()
    return token
async def get_user(token: Annotated[str, Depends(verify_token)]):
    # 根据 token 获取用户信息
    if token == "valid_token":
        return {"username": "john_doe", "user_id": 123}
    else:
        await websocket.close(code=1008)
        raise WebSocketDisconnect()
@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    user: Annotated[dict, Depends(get_user)]
):
    await websocket.accept()
    print(f"Connected user: {user['username']}")
    try:
        while True:
            data = await websocket.receive_text()
            await websocket.send_text(f"Hello {user['username']}, received: {data}")
    except WebSocketDisconnect:
        print(f"User {user['username']} disconnected")

全局依赖（使用 `websocket_route`）

from fastapi.routing import APIRoute
class AuthWebSocketRoute(APIRoute):
    async def get_route_handler(self):
        original_route_handler = super().get_route_handler()
        async def custom_route_handler(websocket: WebSocket):
            # 全局认证逻辑
            token = websocket.headers.get("authorization", "")
            if not token:
                await websocket.close(code=1008)
                return
            # 注入到请求中
            websocket.state.token = token
            websocket.state.user = {"username": "global_user"}
            return await original_route_handler(websocket)
        return custom_route_handler
app = FastAPI()
app.router.route_class = AuthWebSocketRoute
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    user = websocket.state.user
    # ... 处理逻辑

重要注意事项

WebSocket 依赖的特殊性

不能使用 Header、Cookie、Query 等参数装饰器（这些设计用于 HTTP）
需要通过 websocket.headers 或 websocket.query_params 获取数据

错误处理

from fastapi import HTTPException
async def validate_token(websocket: WebSocket) -> str:
    token = websocket.headers.get("authorization", "")
    if not token:
        # 注意：这里不能使用 HTTPException
        await websocket.close(code=1008, reason="Missing token")
        raise WebSocketDisconnect()
    return token

依赖范围

WebSocket 连接的依赖是每次连接时创建
不能使用 @lru_cache 等缓存装饰器（连接生命周期内有效）

性能考虑

# 避免在依赖中做耗时操作
async def slow_dependency():
    # 避免：await some_slow_operation()
    return "result"
# 推荐：使用快速验证
async def fast_dependency(websocket: WebSocket) -> bool:
    # 快速检查
    return websocket.headers.get("authorization") is not None

完整示例：带认证的聊天室

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends
from typing import Annotated
import json
app = FastAPI()
class AuthService:
    async def authenticate(self, websocket: WebSocket) -> dict:
        token = websocket.headers.get("authorization", "")
        # 模拟认证
        if token == "valid_token":
            return {"username": "test_user", "role": "user"}
        await websocket.close(code=1008, reason="Invalid token")
        raise WebSocketDisconnect()
async def get_auth_service():
    return AuthService()
@app.websocket("/chat")
async def chat_endpoint(
    websocket: WebSocket,
    auth_service: Annotated[AuthService, Depends(get_auth_service)]
):
    # 先认证
    user = await auth_service.authenticate(websocket)
    await websocket.accept()
    print(f"User {user['username']} connected")
    try:
        while True:
            data = await websocket.receive_text()
            message = json.loads(data)
            # 处理消息
            response = {
                "username": user["username"],
                "message": message.get("content", ""),
                "timestamp": "2024-01-01T00:00:00Z"
            }
            await websocket.send_json(response)
    except WebSocketDisconnect:
        print(f"User {user['username']} disconnected")

FastAPI WebSocket 支持完整的依赖注入系统，但与 HTTP 路由相比有一些限制：

✅ 支持：