use strawberry dataLoader to solve N+1 query problem

introduce graphql
2026-05-19 00:28:04 +08:00 · 2026-05-19 00:24:45 +08:00
10 changed files with 226 additions and 11 deletions
@@ -5,7 +5,7 @@ prompts:
  preset: "下面的代码来自项目相对路径 {{path}}：\n---\n{{content}}"

 preset:
-  auth: 
+  auth:
    - "app/auth.py"
    - "app/auth/*"

@@ -18,4 +18,5 @@ preset:
  db_model:
    - "app/models/*"
    - "app/db.py"
-    - "app/auth.py"
+    - "app/auth.py"
+    - "app/main.py"
@@ -164,11 +164,11 @@ async def auth_callback(request: Request, code: str, state: str, background_task
        "access_token": jwt_token,
        "refresh_token": refresh_token,
        "token_type": "bearer",
-        "user": {
-            "email": user.email,
-            "name": user.name,
-            "avatar": get_image_url(user.avatar),
-        }
+        # "user": {
+        #     "email": user.email,
+        #     "name": user.name,
+        #     "avatar": get_image_url(user.avatar),
+        # }
    }
    
@router.get("/me")
@@ -0,0 +1,49 @@
+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.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)
+        self.db = SessionLocal()
+        self.request = request
+
+    async def get_current_user(self) -> User | None:
+        # 复用现有的 auth 逻辑
+        try:
+            return await get_current_user(self.request)
+        except Exception:
+            return None
+
+
+async def get_context(request: Request) -> AsyncGenerator[CustomContext, None]:
+    context = CustomContext(request)
+    try:
+        yield context
+    finally:
+        context.db.close()  # 确保请求结束时关闭数据库连接
@@ -0,0 +1,18 @@
+import strawberry
+from app.auth.jwt import create_jwt_token, create_refresh_token
+
+@strawberry.type
+class AuthResponse:
+    access_token: str
+    refresh_token: str
+    token_type: str
+
+@strawberry.type
+class Mutation:
+    @strawberry.mutation
+    def logout(self) -> str:
+        return "ok"
+    
+    # 提示：像 Google OAuth 这种涉及第三方重定向的流程，
+    # 依然保留现有的 RESTful (/login/google 和 /auth/callback) 会更简单。
+    # 登录成功后重定向回前端，前端拿到 token，后续所有数据交互全部走 GraphQL。
@@ -0,0 +1,18 @@
+import strawberry
+from typing import Optional
+from app.graphql.types import UserType
+from app.models import User
+
+@strawberry.type
+class Query:
+    @strawberry.field
+    async def me(self, info) -> Optional[UserType]:
+        # 从 context 获取当前登录用户
+        user = await info.context.get_current_user()
+        # Strawberry 会自动将 SQLAlchemy 实例映射为 UserType
+        return user
+
+    @strawberry.field
+    def get_user_by_handle(self, info, handle: str) -> Optional[UserType]:
+        db = info.context.db
+        return db.query(User).filter(User.handle == handle).first()
@@ -0,0 +1,44 @@
+import strawberry
+from typing import List, Optional
+from datetime import datetime
+from app.utils.format import get_image_url
+
+@strawberry.type
+class UserType:
+    id: strawberry.ID
+    email: str
+    name: Optional[str]
+    handle: Optional[str]
+    avatar: Optional[str]
+    following_count: int
+    follower_count: int
+    created_at: datetime
+
+    # 动态解析字段：例如自动处理头像的绝对路径
+    @strawberry.field
+    def avatar_url(self) -> Optional[str]:
+        return get_image_url(self.avatar) if self.avatar else None
+
+@strawberry.type
+class WorkspaceType:
+    id: strawberry.ID
+    name: str
+    slug: str
+    logo: Optional[str]
+    description: Optional[str]
+    storage_used: int
+    
+    # 支持前端按需加载关联的成员
+    @strawberry.field
+    def members(self, info) -> List["WorkspaceMemberType"]:
+        from app.models import WorkspaceMember
+        db = info.context.db
+        return db.query(WorkspaceMember).filter(WorkspaceMember.workspace_id == self.id).all()
+
+@strawberry.type
+class WorkspaceMemberType:
+    id: strawberry.ID
+    workspace_id: str
+    user_id: Optional[str]
+    role: str
+    status: str
@@ -1,10 +1,31 @@
 from fastapi import FastAPI
-from app.auth import router as auth_router
-from app.models import Base
-from app.db import engine
+from strawberry import Schema
+from strawberry.fastapi import GraphQLRouter

+from app.auth import router as auth_router
+from app.db import engine
+from app.models import Base
+
+# 导入 GraphQL 组件
+from app.graphql.queries import Query
+from app.graphql.mutations import Mutation
+from app.graphql.context import get_context
+
+# 自动建表
 Base.metadata.create_all(bind=engine)

+# 创建 FastAPI 实例
 app = FastAPI()

-app.include_router(auth_router)
+# 组装 GraphQL Schema
+schema = Schema(query=Query, mutation=Mutation)
+
+# 创建 GraphQL 路由并注入上下文
+graphql_router = GraphQLRouter(schema, context_getter=get_context)
+
+# 挂载路由
+# 保留现有的 auth 路由供 Google OAuth 回调使用
+app.include_router(auth_router) 
+
+# 挂载 GraphQL 终点，前端以后只需要请求 /graphql
+app.include_router(graphql_router, prefix="/graphql")
@@ -13,6 +13,7 @@ dependencies = [
    "psycopg2-binary>=2.9.12",
    "python-jose>=3.5.0",
    "sqlalchemy>=2.0.49",
+    "strawberry-graphql[fastapi]>=0.315.5",
    "ulid-py>=1.1.0",
    "uvicorn>=0.47.0",
 ]
@@ -151,6 +151,7 @@ dependencies = [
    { name = "psycopg2-binary" },
    { name = "python-jose" },
    { name = "sqlalchemy" },
+    { name = "strawberry-graphql", extra = ["fastapi"] },
    { name = "ulid-py" },
    { name = "uvicorn" },
 ]
@@ -165,6 +166,7 @@ requires-dist = [
    { name = "psycopg2-binary", specifier = ">=2.9.12" },
    { name = "python-jose", specifier = ">=3.5.0" },
    { name = "sqlalchemy", specifier = ">=2.0.49" },
+    { name = "strawberry-graphql", extras = ["fastapi"], specifier = ">=0.315.5" },
    { name = "ulid-py", specifier = ">=1.1.0" },
    { name = "uvicorn", specifier = ">=0.47.0" },
 ]
@@ -237,6 +239,18 @@ version = "1.7"
 source = { registry = "https://pypi.org/simple" }
 sdist = { url = "https://files.pythonhosted.org/packages/6b/b0/e595ce2a2527e169c3bcd6c33d2473c1918e0b7f6826a043ca1245dd4e5b/crcmod-1.7.tar.gz", hash = "sha256:dc7051a0db5f2bd48665a990d3ec1cc305a466a77358ca4492826f41f283601e", size = 89670, upload-time = "2010-06-27T14:35:29.538Z" }

+[[package]]
+name = "cross-web"
+version = "0.6.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ad/83/b5ef04565acc065387dda3a4fbf0c4cfb6bab805c81b66b2bc5b5ac9a282/cross_web-0.6.0.tar.gz", hash = "sha256:ae90570802615365ca1a781117b43bfd0d6cd3bf611649d24c3a206a82a693c9", size = 331315, upload-time = "2026-04-13T14:29:12.718Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/35/a2/dab06d9b80cb76c700883186a9a2e6fd103342c9b4def4d88f5787796e17/cross_web-0.6.0-py3-none-any.whl", hash = "sha256:bdebf0c08d02f3a48cf67b6904d3a6d8fd8cab2cd905592ab96ab00b259cd582", size = 24820, upload-time = "2026-04-13T14:29:11.198Z" },
+]
+
 [[package]]
 name = "cryptography"
 version = "48.0.0"
@@ -318,6 +332,15 @@ wheels = [
    { url = "https://files.pythonhosted.org/packages/5a/ff/2e4eca3ade2c22fe1dea7043b8ee9dabe47753349eb1b56a202de8af6349/fastapi-0.136.1-py3-none-any.whl", hash = "sha256:a6e9d7eeada96c93a4d69cb03836b44fa34e2854accb7244a1ece36cd4781c3f", size = 117683, upload-time = "2026-04-23T16:49:42.437Z" },
 ]

+[[package]]
+name = "graphql-core"
+version = "3.2.8"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/68/c5/36aa96205c3ecbb3d34c7c24189e4553c7ca2ebc7e1dd07432339b980272/graphql_core-3.2.8.tar.gz", hash = "sha256:015457da5d996c924ddf57a43f4e959b0b94fb695b85ed4c29446e508ed65cf3", size = 513181, upload-time = "2026-03-05T19:55:37.332Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/86/41/cb887d9afc5dabd78feefe6ccbaf83ff423c206a7a1b7aeeac05120b2125/graphql_core-3.2.8-py3-none-any.whl", hash = "sha256:cbee07bee1b3ed5e531723685369039f32ff815ef60166686e0162f540f1520c", size = 207349, upload-time = "2026-03-05T19:55:35.911Z" },
+]
+
 [[package]]
 name = "greenlet"
 version = "3.5.0"
@@ -452,6 +475,15 @@ dependencies = [
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/df/b5/f2cb1950dda46ac2284d6c950489fdacd0e743c2d79a347924d3cc44b86f/oss2-2.19.1.tar.gz", hash = "sha256:a8ab9ee7eb99e88a7e1382edc6ea641d219d585a7e074e3776e9dec9473e59c1", size = 298845, upload-time = "2024-10-25T11:37:46.638Z" }

+[[package]]
+name = "packaging"
+version = "26.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" },
+]
+
 [[package]]
 name = "psycopg2-binary"
 version = "2.9.12"
@@ -590,6 +622,15 @@ wheels = [
    { url = "https://files.pythonhosted.org/packages/d9/c3/0bd11992072e6a1c513b16500a5d07f91a24017c5909b02c72c62d7ad024/python_jose-3.5.0-py2.py3-none-any.whl", hash = "sha256:abd1202f23d34dfad2c3d28cb8617b90acf34132c7afd60abd0b0b7d3cb55771", size = 34624, upload-time = "2025-05-28T17:31:52.802Z" },
 ]

+[[package]]
+name = "python-multipart"
+version = "0.0.29"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/4e/fe/70bd71a6738b09a0bdf6480ca6436b167469ca4578b2a0efbe390b4b0e70/python_multipart-0.0.29.tar.gz", hash = "sha256:643e93849196645e2dbdd81a0f8829a23123ad7f797a84a364c6fb3563f18904", size = 45678, upload-time = "2026-05-17T17:29:47.654Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8f/cb/769cfc37177252872a45a71f3fbdde9d51b471a3f3c14bfe95dde3407386/python_multipart-0.0.29-py3-none-any.whl", hash = "sha256:2ddcc971cef266225f54f552d8fa10bcfbb1f14446caec199060daac59ff2d69", size = 29640, upload-time = "2026-05-17T17:29:45.69Z" },
+]
+
 [[package]]
 name = "requests"
 version = "2.34.2"
@@ -676,6 +717,28 @@ wheels = [
    { url = "https://files.pythonhosted.org/packages/0b/c9/584bc9651441b4ba60cc4d557d8a547b5aff901af35bda3a4ee30c819b82/starlette-1.0.0-py3-none-any.whl", hash = "sha256:d3ec55e0bb321692d275455ddfd3df75fff145d009685eb40dc91fc66b03d38b", size = 72651, upload-time = "2026-03-22T18:29:45.111Z" },
 ]

+[[package]]
+name = "strawberry-graphql"
+version = "0.315.5"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "cross-web" },
+    { name = "graphql-core" },
+    { name = "packaging" },
+    { name = "python-dateutil" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/18/9b/101940ee899959d4d220cc8beb4f41bce9c58830d1872df06f35cf92c457/strawberry_graphql-0.315.5.tar.gz", hash = "sha256:29a2f04479aba29f9f30ecdfce1ef9ee04acee3c434a4f6019249474cd649492", size = 222724, upload-time = "2026-05-14T10:46:16.222Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7e/19/e389140b3b50faf803b4817ae5a8c3c2d6737f062f4a0fe6778e929a405d/strawberry_graphql-0.315.5-py3-none-any.whl", hash = "sha256:073bc818a5f55951a9a6fbab40bbfa07c418d2c4151cc2aab24399bd14d9a51a", size = 325062, upload-time = "2026-05-14T10:46:18.672Z" },
+]
+
+[package.optional-dependencies]
+fastapi = [
+    { name = "fastapi" },
+    { name = "python-multipart" },
+]
+
 [[package]]
 name = "typing-extensions"
 version = "4.15.0"
Author	SHA1	Message	Date
Zayden-Jung	a83befe109	use strawberry dataLoader to solve N+1 query problem	2026-05-19 00:28:04 +08:00
Zayden-Jung	9c55cc675e	introduce graphql	2026-05-19 00:24:45 +08:00