Celery动态分配会话/复用连接：多数据库查询连接复用问题求助

核心思路 #

Celery的Task基类是单例模式，装饰器中绑定的参数会在任务注册时固定，无法动态切换数据库。因此我们需要：

1. 实现一个按数据库标识缓存连接池的DBContext，避免重复初始化数据库引擎
2. 用上下文管理器自动处理会话的创建与释放，确保任务完成后会话正确关闭
3. 编写通用任务函数，动态接收数据库配置参数，无需为每个数据库单独编写任务

1. 实现可复用的DBContext #

这个类会缓存每个数据库的引擎和会话工厂，在Worker进程生命周期内复用连接池，同时为每个任务提供独立的会话：

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session

class DBContext:
    # 类变量：缓存已初始化的数据库引擎和会话工厂
    _engines = {}
    _session_factories = {}

    @classmethod
    def _get_db_key(cls, db_config):
        # 生成数据库唯一标识（用于缓存）
        return f"{db_config['drivername']}://{db_config['username']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}"

    @classmethod
    def get_engine(cls, db_config):
        db_key = cls._get_db_key(db_config)
        if db_key not in cls._engines:
            # 初始化数据库引擎，配置连接池
            cls._engines[db_key] = create_engine(
                db_key,
                pool_size=10,       # 常驻连接数
                max_overflow=20,    # 额外临时连接数
                pool_recycle=3600   # 自动回收闲置连接
            )
            cls._session_factories[db_key] = sessionmaker(bind=cls._engines[db_key])
        return cls._engines[db_key]

    @classmethod
    def get_session(cls, db_config):
        # 返回上下文管理器，自动创建/关闭会话
        db_key = cls._get_db_key(db_config)
        session_factory = cls._session_factories[db_key]
        # 使用scoped_session确保线程安全（Celery Worker默认是多进程，每个进程内线程安全）
        session = scoped_session(session_factory)
        try:
            yield session
        finally:
            # 任务完成后销毁会话，释放连接回池
            session.remove()

2. 编写通用Celery任务 #

直接在任务中动态接收数据库配置，通过DBContext获取会话，无需依赖固定的Task基类：

from sqlalchemy import text

@app.task(bind=True, name='fetch_data')
def fetch_data(self, db_config, sql):
    # 用上下文管理器自动管理会话生命周期
    with DBContext.get_session(db_config) as session:
        # 执行原生SQL查询（如果是ORM查询，直接传入模型类即可）
        result = session.execute(text(sql)).fetchall()
        # 转换结果为易处理的字典格式
        return [dict(row) for row in result]

3. 调用任务示例 #

只需传入目标数据库的配置字典和SQL语句即可：

# 数据库A的配置
db_config_a = {
    "drivername": "postgresql",
    "username": "user_a",
    "password": "pass_a",
    "host": "db-a.example.com",
    "port": 5432,
    "database": "db_a"
}

# 数据库B的配置
db_config_b = {
    "drivername": "mysql+pymysql",
    "username": "user_b",
    "password": "pass_b",
    "host": "db-b.example.com",
    "port": 3306,
    "database": "db_b"
}

# 向不同数据库发送任务
fetch_data.delay(db_config_a, "SELECT * FROM table_a")
fetch_data.delay(db_config_b, "SELECT * FROM table_b")

关键说明 #

• 连接池复用：DBContext的_engines是类变量，会在每个Celery Worker进程内缓存，避免重复创建数据库连接，提升性能。
• 会话隔离：每个任务都会获取独立的会话，任务完成后自动销毁，避免数据污染和连接泄漏。
• 完全通用：无需为每个数据库编写单独的任务函数，所有数据库查询都可以通过这一个任务处理。

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