Django生产环境下如何实现类似assertNumQueries的SQL查询次数统计上下文管理器

嘿，这个需求我之前也碰到过！Django自带的connection.queries确实只在DEBUG模式下生效，生产环境根本没法直接用，但我们完全可以自己实现一个类似测试用例里self.assertNumQueries()的上下文管理器，下面分享两种靠谱的实现方式：

方式一：极简版通用统计 #

这个方法利用Django数据库连接的execute_wrapperAPI，直接给查询执行过程加个计数器，不依赖DEBUG模式，生产环境完全可用：

from contextlib import contextmanager
from django.db import connection

@contextmanager
def count_sql_queries():
    query_count = 0

    def increment_count(execute, sql, params, many, context):
        nonlocal query_count
        query_count += 1
        # 继续执行原本的查询逻辑
        return execute(sql, params, many, context)

    # 给当前数据库连接绑定计数包装器，上下文结束后自动清理
    with connection.execute_wrapper(increment_count):
        yield query_count

然后就可以完全按照你想要的方式使用：

def my_view(request):
    with count_sql_queries() as nr_queries:
        response = business_logic()
        logger.debug(f'There were {nr_queries} queries made')
    return response

为啥这个能在生产环境跑？ #

• 它用的是Django数据库连接的通用API，和DEBUG模式无关
• 只会在上下文管理器的范围内统计，不会干扰其他请求的查询计数
• 支持所有Django兼容的数据库后端，不用操心适配问题

方式二：区分读/写查询的精细统计 #

如果需要区分SELECT（读操作）和INSERT/UPDATE/DELETE（写操作）的查询次数，可以稍微改造一下：

from contextlib import contextmanager
from django.db import connection
import re

@contextmanager
def count_sql_queries():
    query_counts = {'total': 0, 'read': 0, 'write': 0}
    # 用正则匹配SELECT开头的查询（忽略大小写和前置空格）
    read_query_pattern = re.compile(r'^\s*SELECT', re.IGNORECASE)

    def increment_count(execute, sql, params, many, context):
        nonlocal query_counts
        # 判断查询类型
        if read_query_pattern.match(sql):
            query_counts['read'] += 1
        else:
            query_counts['write'] += 1
        query_counts['total'] += 1
        return execute(sql, params, many, context)

    with connection.execute_wrapper(increment_count):
        yield query_counts

使用时可以拿到更详细的统计数据：

def my_view(request):
    with count_sql_queries() as counts:
        response = business_logic()
        logger.debug(f'Total queries: {counts["total"]}, Reads: {counts["read"]}, Writes: {counts["write"]}')
    return response

生产环境使用注意事项 #

• 如果你的应用用了多数据库，这个管理器默认只统计默认数据库的查询。要统计其他库的话，把connection换成connections['你的数据库别名']就行
• 这个统计会包含所有Django ORM发起的查询，包括隐式的关联查询、缓存失效后的重新查询等，和测试用例里assertNumQueries的统计逻辑完全一致
• 虽然这个包装器的性能开销极小，但高并发场景下建议只在排查性能问题的视图中临时添加，或者通过日志级别控制输出（比如只在DEBUG级别的日志里打印），避免不必要的资源消耗

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