FastAPI：如何在自定义Middleware的dispatch函数中传入get_db依赖？

由于BaseHTTPMiddleware的dispatch方法无法直接使用FastAPI的依赖注入系统，你可以通过以下两种方式实现数据库查询：

方案1：直接创建数据库会话 #

如果你的get_db是基于SQLAlchemy的典型实现，可直接实例化会话类获取连接，用完手动关闭：

先看标准的SQLAlchemy数据库配置（假设你用的是这个）：

from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

SQLALCHEMY_DATABASE_URL = "sqlite:///./test.db"

engine = create_engine(
    SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

修改Middleware代码：

from fastapi import Request, BaseHTTPMiddleware, JSONResponse
from sqlalchemy.orm import Session

class MyMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        access_token = request.headers.get("access-token")
        if not access_token:
            return JSONResponse(status_code=401, content={"errors": "access_token not found"})
        
        # 直接初始化数据库会话
        db: Session = SessionLocal()
        try:
            # 这里写你的数据库查询逻辑
            # 示例：user = db.query(User).filter(User.access_token == access_token).first()
            pass
        finally:
            # 必须关闭会话，防止连接泄漏
            db.close()
        
        response = await call_next(request)
        return response

方案2：将get_db依赖注入Middleware构造函数 #

这种方式更贴合FastAPI的依赖设计，把get_db传入Middleware的构造函数，在dispatch中调用获取会话：

from fastapi import Request, BaseHTTPMiddleware, JSONResponse, Depends
from sqlalchemy.orm import Session

class MyMiddleware(BaseHTTPMiddleware):
    def __init__(self, app, db_dependency=Depends(get_db)):
        super().__init__(app)
        self.db_dependency = db_dependency

    async def dispatch(self, request: Request, call_next):
        access_token = request.headers.get("access-token")
        if not access_token:
            return JSONResponse(status_code=401, content={"errors": "access_token not found"})
        
        # 通过依赖获取数据库会话
        db: Session = next(self.db_dependency())
        try:
            # 执行你的数据库操作
            # 示例：user = db.query(User).filter(User.access_token == access_token).first()
            pass
        finally:
            db.close()
        
        response = await call_next(request)
        return response

注册Middleware时直接添加即可：

app = FastAPI()
app.add_middleware(MyMiddleware)

异步数据库注意事项 #

如果使用异步数据库驱动（如asyncpg），需改用异步会话，并配合async with语法确保会话正确关闭：

# 异步会话配置示例
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker

async_engine = create_async_engine("postgresql+asyncpg://user:pass@localhost/db")
AsyncSessionLocal = sessionmaker(async_engine, class_=AsyncSession, expire_on_commit=False)

# Middleware中的异步查询
async with AsyncSessionLocal() as db:
    result = await db.execute(select(User).where(User.access_token == access_token))
    user = result.scalar_one_or_none()

内容的提问来源于stack exchange，提问作者Fahad Yamin