如何在FastAPI中实现基于请求动态连接多异步数据库？

核心思路 #

根据请求体中user_id的范围映射到对应数据库连接，在请求处理阶段动态选择异步引擎并创建会话，确保每个请求的会话绑定到正确的数据库，请求结束后自动回收资源。

实现步骤 #

1. 预初始化异步引擎与会话工厂 #

先定义所有数据库的连接配置，批量创建异步引擎和对应的会话工厂：

from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker

# 数据库连接配置
DB_MAP = {
    "user1-10k": "mysql+aiomysql://username:password@host/db_user_1_10k",
    "user10k-50k": "mysql+aiomysql://username:password@host/db_user_10k_50k",
    # 可扩展更多分片库
}

# 存储引擎和会话工厂的字典
async_engines: dict[str, AsyncEngine] = {}
async_session_factories: dict[str, sessionmaker] = {}

# 初始化所有引擎与会话工厂
for db_key, db_url in DB_MAP.items():
    engine = create_async_engine(db_url, echo=False)
    async_engines[db_key] = engine
    async_session_factories[db_key] = sessionmaker(
        bind=engine, class_=AsyncSession, expire_on_commit=False
    )

2. 编写数据库路由规则 #

根据user_id范围返回对应的数据库标识：

def get_db_key(user_id: int) -> str:
    if 1 <= user_id <= 10000:
        return "user1-10k"
    elif 10001 <= user_id <= 50000:
        return "user10k-50k"
    else:
        raise ValueError("user_id不在支持的范围内")

3. 依赖注入动态获取会话 #

通过FastAPI依赖注入，在请求阶段解析请求体、选择数据库并返回会话：

from fastapi import FastAPI, Depends, HTTPException
from pydantic import BaseModel

app = FastAPI()

# 请求体模型
class UserRequest(BaseModel):
    user_id: int
    status: str

# 动态会话依赖
async def get_db_session(req: UserRequest) -> AsyncSession:
    try:
        db_key = get_db_key(req.user_id)
        session_factory = async_session_factories[db_key]
        async with session_factory() as session:
            yield session
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

# 示例接口
@app.post("/handle-user")
async def handle_user(
    req: UserRequest,
    session: AsyncSession = Depends(get_db_session)
):
    # 这里执行数据库操作，比如：
    # result = await session.execute(select(User).where(User.id == req.user_id))
    return {
        "msg": f"已连接数据库 {get_db_key(req.user_id)}",
        "user_id": req.user_id,
        "status": req.status
    }

4. 灵活场景适配（可选） #

如果不需要强校验请求体，可直接解析原始请求体获取user_id：

from fastapi import Request

async def get_db_session(request: Request) -> AsyncSession:
    try:
        body = await request.json()
        user_id = int(body.get("user_id"))
        db_key = get_db_key(user_id)
        session_factory = async_session_factories[db_key]
        async with session_factory() as session:
            yield session
    except (TypeError, ValueError) as e:
        raise HTTPException(status_code=400, detail="无效的user_id参数")

关于官方多引擎方案无效的说明 #

SQLAlchemy官方的分片路由（如ShardedSession）是针对查询级的自动路由，依赖查询条件自动匹配分片，而我们的场景是请求级的数据库选择——需要在请求入口就确定数据库，因此直接通过依赖注入绑定对应引擎的会话是更贴合需求的方案。

内容的提问来源于stack exchange，提问作者Henryk ll