如何在Python SQLAlchemy中程序化测量查询复杂度并管控查询次数？

完全可以通过SQLAlchemy的事件监听机制实现这个需求，核心思路是统计执行的SQL查询次数，超过预设阈值时抛出错误。下面是具体实现方案：

1. 实现查询计数的上下文管理器 #

利用SQLAlchemy的after_execute事件，在查询执行后累加计数，退出上下文时检查是否超过阈值：

from sqlalchemy import event
from sqlalchemy.orm import Session

class QueryLimitExceededError(Exception):
    pass

class QueryCounter:
    def __init__(self, session: Session, max_queries: int):
        self.session = session
        self.max_queries = max_queries
        self.query_count = 0
        self.listener = None

    def _count_query(self, conn, clauseelement, multiparams, params, execution_options, result):
        self.query_count += 1
        if self.query_count > self.max_queries:
            raise QueryLimitExceededError(f"查询次数超过阈值：{self.query_count} > {self.max_queries}")

    def __enter__(self):
        self.listener = event.listen(self.session.bind, "after_execute", self._count_query)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        event.remove(self.session.bind, "after_execute", self._count_query)

2. 在单元测试中使用计数器 #

比如测试查询单个Parent时，预期最多1次查询（预加载children的情况），如果没做预加载，查询次数会是2次（1次查Parent，1次查Children），此时会触发错误：

def test_parent_query_without_n_plus_1():
    session = create_test_session()
    # 插入测试数据：1个Parent带2个Children
    parent = ParentModel()
    session.add(parent)
    session.add(ChildModel(parent_id=parent.id))
    session.add(ChildModel(parent_id=parent.id))
    session.commit()

    # 预期最多1次查询（预加载Children的情况）
    with QueryCounter(session, max_queries=1):
        # 模拟GraphQL查询请求
        query = """
            query {
                parent(id: 1) {
                    id
                    children {
                        id
                    }
                }
            }
        """
        result = schema.execute_sync(query, context_value={"session": session})
        assert not result.errors
        assert len(result.data["parent"]["children"]) == 2

如果代码里没做预加载，这段测试会抛出QueryLimitExceededError，直接暴露N+1问题。

3. 集成到Strawberry请求流程（可选） #

如果想在开发环境实时监控，也可以把计数器集成到Strawberry的上下文处理中，每个请求自动启动计数：

from strawberry.fastapi import GraphQLRouter

def create_graphql_router(schema):
    async def get_context():
        session = create_db_session()
        counter = QueryCounter(session, max_queries=10)  # 按需设置阈值
        counter.__enter__()
        try:
            yield {"session": session, "counter": counter}
        finally:
            counter.__exit__(None, None, None)
            session.close()

    return GraphQLRouter(schema, context_getter=get_context)

这样每次请求都会自动监控查询次数，超过阈值直接报错，在开发阶段就能及时发现N+1问题。

4. 修复示例中的N+1问题 #

最后补充原代码的修复方案，使用SQLAlchemy的joinedload预加载关联模型，避免多次查询：

from sqlalchemy.orm import joinedload

@strawberry.type
class Query:
    @strawberry.field
    def parent(self, info, id: int) -> ParentGQLSchema | None:
        model = session.query(ParentModel)\
            .options(joinedload(ParentModel.children))\  # 预加载children
            .filter(ParentModel.id == id)\
            .first()
        if not model:
            return None
        return ParentGQLSchema.from_model(model)

# 修改ParentGQLSchema的children字段，直接用预加载的模型
@strawberry.type
class ParentGQLSchema:
    id: int
    _children: list[ChildModel]  # 存储预加载的模型

    @classmethod
    def from_model(cls, model: ParentModel):
        return cls(
            id=model.id,
            _children=model.children
        )

    @strawberry.field
    def children(
        self, info, page: int = 1, limit: int = 20
    ) -> list[ChildGQLSchema]:
        # 直接用预加载的_children做分页处理
        start = (page - 1) * limit
        end = start + limit
        return [ChildGQLSchema.from_model(child) for child in self._children[start:end]]

这样查询单个Parent时只会执行1次SQL查询，彻底解决N+1问题。

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