diff --git a/app/graphql/context.py b/app/graphql/context.py index 20d7b5c..45c37b9 100644 --- a/app/graphql/context.py +++ b/app/graphql/context.py @@ -1,10 +1,32 @@ +from app.models.user import User +from app.auth.deps import get_current_user from fastapi import Request 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.auth.deps import get_current_user -from app.models import User +from app.models import WorkspaceMember +from app.graphql.types import WorkspaceMemberType 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): def __init__(self, request: Request): super().__init__(request) @@ -18,6 +40,7 @@ class CustomContext(BaseContext): except Exception: return None + async def get_context(request: Request) -> AsyncGenerator[CustomContext, None]: context = CustomContext(request) try: