如何在Django中为MariaDB/MySQL查询设置超时并执行启发式查询

核心思路 #

MySQL与MariaDB不支持PostgreSQL那样的会话级SET LOCAL超时设置，且你无法直接修改生成的SQL前缀，因此我们通过临时设置会话级超时参数的方式实现超时控制，执行完成后再恢复默认值，同时针对两者的参数差异做适配：

• MariaDB使用max_statement_time，单位为秒
• MySQL 5.7+使用max_execution_time，单位为毫秒

实现代码 #

from django.db import connection, OperationalError
from django.db.transaction import atomic

def execute_with_timeout_fallback(orm_query, heuristic_query, timeout_seconds=1):
    # 根据数据库类型匹配超时参数
    db_vendor = connection.vendor
    if db_vendor == 'mysql':
        timeout_param = 'max_execution_time'
        timeout_value = timeout_seconds * 1000  # 转毫秒
    else:
        timeout_param = 'max_statement_time'
        timeout_value = timeout_seconds  # 直接用秒

    with atomic():
        with connection.cursor() as cursor:
            try:
                # 临时设置会话级超时
                cursor.execute(f"SET SESSION {timeout_param} = {timeout_value};")
                # 处理Django ORM查询：转换为带参数的SQL
                sql, params = orm_query.query.sql_with_params()
                cursor.execute(sql, params)
                # 根据业务需求返回结果，示例为返回所有结果
                return cursor.fetchall()
            except OperationalError as e:
                # 仅捕获超时错误：MySQL错误码3024，MariaDB错误码1969
                if e.args[0] in (3024, 1969):
                    # 执行启发式查询
                    cursor.execute(heuristic_query)
                    return int(cursor.fetchone()[0])
                # 非超时类操作错误，重新抛出
                raise
            finally:
                # 恢复会话超时为默认值（0表示无限制）
                cursor.execute(f"SET SESSION {timeout_param} = 0;")

关键细节说明 #

• 会话参数恢复：必须在finally块重置超时参数，避免影响后续数据库请求
• 错误精准捕获：仅针对超时特定错误码触发降级，防止误处理其他操作异常
• ORM查询安全处理：通过query.sql_with_params()获取SQL和参数，避免手动拼接导致的注入风险
• 事务上下文：atomic()保证操作原子性，若为只读场景可调整隔离级别或移除

备选方案（若会话级设置不适配连接池） #

如果项目使用数据库连接池，会话级参数可能影响其他并发请求，可直接用SET STATEMENT语法包装SQL（需能获取原始SQL）：

def wrap_query_with_timeout(sql, timeout_seconds, is_mariadb):
    if is_mariadb:
        return f"SET STATEMENT max_statement_time={timeout_seconds} FOR {sql}"
    else:
        return f"SET STATEMENT max_execution_time={timeout_seconds*1000} FOR {sql}"

# 使用示例
sql, params = orm_query.query.sql_with_params()
wrapped_sql = wrap_query_with_timeout(sql, timeout_seconds, db_vendor != 'mysql')
cursor.execute(wrapped_sql, params)

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