如何在没有全局变量的芹菜任务中设置 sqlalchemy 会话

Question

总结：我想在 celery 任务中使用 sqlalchemy 会话，而不需要包含该会话的全局变量。

我在一个带有 celery 任务的项目中使用 SQLAlchemy，我正在使用

目前，我在 celery 应用程序设置 (celery.py) 中定义了一个全局变量“会话”，并带有一个工作信号来设置它。

session = scoped_session(sessionmaker())

@celeryd_init.connect
def configure_workers(sender=None, conf=None, **kwargs):
    # load the application configuration
    # db_uri = conf['db_uri']
    engine = create_engine(db_uri)
    session.configure(bind=engine)

在定义任务的模块中，我只需导入“会话”并使用它。任务是使用自定义类定义的，该类在返回后关闭会话：

class DBTask(Task):
    def after_return(self, *args, **kwargs):
        session.remove()

但效果很好：当使用 CELERY_ALWAYS_EAGER=True 进行单元测试时，不会配置会话。到目前为止，我发现的唯一解决方案是在单元测试中运行任务时模拟“会话”变量：

with mock.patch('celerymodule.tasks.session', self.session):
    do_something.delay(...)

虽然它有效，但我不想这样做。

有什么方法可以设置一个不是全局变量的会话，它既适用于正常的异步行为，也适用于没有 CELERY_ALWAYS_EAGER=True 的工作人员？

Answer 1

关于custom task classes的官方文档中的答案就在我的眼皮底下。

我修改了用于访问数据库的任务的自定义任务类：

class DBTask(Task):
    _session = None

    def after_return(self, *args, **kwargs):
        if self._session is not None:
            self._session.remove()

    @property
    def session(self):
        if self._session is None:
            _, self._session = _get_engine_session(self.conf['db_uri'],
                                                   verbose=False)

        return self._session

我这样定义我的任务：

@app.task(base=DBTask, bind=True)
def do_stuff_with_db(self, conf, some_arg):
    self.conf = conf
    thing = self.session.query(Thing).filter_by(arg=some_arg).first()

这样，SQLAlchemy 会话只会为每个 celery 工作进程创建一次，我不需要任何全局变量。

这解决了我的单元测试问题，因为 SQLAlchemy 会话设置现在独立于 celery 工人。