Django Celery长时任务Postgres数据库连接丢失的优化方案咨询

针对长时Celery任务丢失Postgres连接的问题，以下是几个更严谨且易测试的优化方案：

1. 使用Django内置连接健康检查替代手动重复连接 #

直接调用connection.connect()会强制重建连接，容易在单元测试的事务回滚阶段引发错误。改用Django提供的connection.ensure_connection()方法，它会先检查当前连接是否有效，仅在连接断开时才重建，既解决连接丢失问题，又不干扰测试环境的事务管理。

同时建议缩小事务范围，避免长事务长时间占用连接，进一步降低连接丢失概率，出错时仅回滚当前批次数据：

from django.db import connection, transaction

@app.task(name='my_name_of_the_task')
def my_long_running_task(params):
    object_list = self.get_object_list_from_params(params)
    for batch in batches(object_list):
        # 确保连接有效，无效则自动重建
        connection.ensure_connection()
        with transaction.atomic():
            self.some_calculations()
            MyObject.objects.bulk_create(objs=batch)

2. 封装特定异常的重试逻辑 #

不要捕获所有Exception，仅针对Postgres连接相关异常（如psycopg2.OperationalError）处理，避免掩盖其他业务错误。可以封装一个复用性强的重试装饰器：

import psycopg2
from functools import wraps
from django.db import connection

def retry_on_db_disconnect(max_retries=2):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            retries = 0
            while retries <= max_retries:
                try:
                    return func(*args, **kwargs)
                except psycopg2.OperationalError:
                    # 重置连接后重试
                    connection.close()
                    connection.ensure_connection()
                    retries += 1
            raise Exception("数据库连接重试次数耗尽")
        return wrapper
    return decorator

在任务中使用装饰器：

@app.task(name='my_name_of_the_task')
def my_long_running_task(params):
    object_list = self.get_object_list_from_params(params)
    for batch in batches(object_list):
        @retry_on_db_disconnect(max_retries=2)
        def process_batch():
            with transaction.atomic():
                self.some_calculations()
                MyObject.objects.bulk_create(objs=batch)
        
        process_batch()

3. 结合Celery自带的任务重试机制 #

如果本地重试无法解决严重连接故障，可以利用Celery的任务重试功能，设置合理的重试间隔和次数，同时确保任务具备幂等性（比如给MyObject添加唯一约束，或在获取数据时过滤已创建的对象）：

import psycopg2
from celery.exceptions import Retry

@app.task(
    name='my_name_of_the_task',
    autoretry_for=(psycopg2.OperationalError,),
    retry_backoff=3,
    retry_kwargs={'max_retries': 3}
)
def my_long_running_task(params):
    object_list = self.get_object_list_from_params(params)
    for batch in batches(object_list):
        with transaction.atomic():
            self.some_calculations()
            MyObject.objects.bulk_create(objs=batch)

retry_backoff设置重试间隔指数增长（3秒、6秒、12秒），避免短时间内频繁重试给数据库造成压力。

额外注意事项 #

• 避免长事务：长事务不仅易导致连接丢失，还会占用数据库资源、增加锁竞争，尽量将事务拆分为小范围操作。
• 单元测试验证：测试时可使用Django的TestCase自动管理事务，用mock库模拟psycopg2.OperationalError，验证重试逻辑是否生效。

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