Fastapi自定义中间件动态处理多数据库连接

背景

        自己的fastapi服务,需要连接不同的数据库实例支持多维度业务.查阅了官方文档,也在网上找了很久,也没有找到直接可以采用的方案,既然找不到,那就自己来

        首先,连接多数据库,配置多套数据库连接,实现不同的db实例对象即可实现.然后在不同的router中分别调用.low否?那能不能有一个公共的get_db.供不同的router依赖注入,来完成动态的创建dbSession对象呢?

方案

        我们需要根据定义的路由去区分要连接的数据库,那么如何获取请求的路由呢?是不是可以用中间件来拦截请求的url,然后解析出路由名称呢? 获取到了路由,我们就可以通过路由与数据库连接的映射,来连接相应的库进行处理

实现

        首先上代码        

class DatabaseSelector(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)

    async def dispatch(self, request: Request, call_next):
        url_param = str(request.url).split("/")
        if "api" in str(request.url):
            url_param.remove("api")
        route_name = url_param[3]
        if route_name not in list(routeMap.keys()):
            SQLALCHEMY_DATABASE_URI = routeMap["base_mysql"]
        else:
            SQLALCHEMY_DATABASE_URI = routeMap[route_name]
        os.environ["SQLALCHEMY_DATABASE_URI"] = SQLALCHEMY_DATABASE_URI
        response = await call_next(request)
        return response

自定义的中间件继承了 BaseHTTPMiddleware, 重写了dispatch方法处理我们的需求

在请求中,通过request.url 获取请求的完整url,然后截取路由名称. 处理"api"是为了兼容本地和正式环境的请求url的包装

通过映射拿到了目标数据库连接,将变量放到临时环境中,供实例化Session时获取

# 声明SQLAlchemy对象类型
Base = declarative_base()


def get_db():
    SQLALCHEMY_DATABASE_URI = os.getenv('SQLALCHEMY_DATABASE_URI')
    engine = create_engine(SQLALCHEMY_DATABASE_URI, echo=True, pool_size=50)
    SessionLocal = sessionmaker(bind=engine)
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

在config中配置数据库连接信息,以及映射配置

# 数据库映射
# base_mysql 是平台主连接 其他的连接的key要与自身服务的router名称保持一致, 如 show_data
routeMap = {
    "base_mysql": 'mysql+pymysql://{}:{}@{}:{}/{}'.format(*BaseMysql.values()),
    "show_data": 'mysql+pymysql://{}:{}@{}:{}/{}'.format(*RequestMysql.values()),
    "else_data": "xxxxxxx",
}

最终Depends依赖注入get_db,对不同的目标库建立连接,获取Session对象

@router.post("/dashboard_select_project", response_model=ResPonseModel)
async def dashboard_select_project(selects: DashboardSelect, db=Depends(get_db)):

最后

        第一篇文章,不是很喜欢码字,但是确实是找现成的资料太难了,可能是fastapi应用的还不是很多,资料少,也可能是我的嗅探姿势不对...找到的都是官方文档,以及一些没人回答的问题贴.不管怎样,既然搞了.就分享给大家,希望和我一样不停的搜各种关键字组合,却没有满意答案的伙伴们,能够提供给大家一些帮助.