use strawberry dataLoader to solve N+1 query problem

This commit is contained in:
2026-05-19 00:28:04 +08:00
parent 9c55cc675e
commit a83befe109
+25 -2
View File
@@ -1,10 +1,32 @@
from app.models.user import User
from app.auth.deps import get_current_user
from fastapi import Request from fastapi import Request
from strawberry.fastapi import BaseContext from strawberry.fastapi import BaseContext
from strawberry.dataloader import DataLoader
from typing import List
from collections import defaultdict
from app.db import SessionLocal from app.db import SessionLocal
from app.auth.deps import get_current_user from app.models import WorkspaceMember
from app.models import User from app.graphql.types import WorkspaceMemberType
from typing import AsyncGenerator from typing import AsyncGenerator
# 定义批量查询函数:接收一组 workspace ID 列表,返回相同长度的结果列表
async def load_members_by_workspace_ids(workspace_ids: List[str]) -> List[List[WorkspaceMember]]:
db = SessionLocal()
try:
# 用一条 IN 语句查出所有相关成员
members = db.query(WorkspaceMember).filter(WorkspaceMember.workspace_id.in_(workspace_ids)).all()
# 将结果归类到对应的 workspace ID
member_map = defaultdict(list)
for member in members:
member_map[member.workspace_id].append(member)
# 严格按照传入的 workspace ID 顺序返回结果列表
return [member_map[w_id] for w_id in workspace_ids]
finally:
db.close()
class CustomContext(BaseContext): class CustomContext(BaseContext):
def __init__(self, request: Request): def __init__(self, request: Request):
super().__init__(request) super().__init__(request)
@@ -18,6 +40,7 @@ class CustomContext(BaseContext):
except Exception: except Exception:
return None return None
async def get_context(request: Request) -> AsyncGenerator[CustomContext, None]: async def get_context(request: Request) -> AsyncGenerator[CustomContext, None]:
context = CustomContext(request) context = CustomContext(request)
try: try: