Strawberry GraphQL 中的 DataLoader 使用指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00178/article/details/148527493

Strawberry GraphQL 中的 DataLoader 使用指南

strawberry A GraphQL library for Python that leverages type annotations 🍓 项目地址: https://gitcode.com/gh_mirrors/st/strawberry

前言

在现代 GraphQL 应用中，N+1 查询问题是一个常见的性能瓶颈。Strawberry GraphQL 框架内置了 DataLoader 工具，它能有效解决这个问题。本文将深入介绍如何在 Strawberry 中使用 DataLoader 来优化数据加载性能。

DataLoader 基础概念

DataLoader 是一个批处理和缓存工具，它通过以下两种机制提升性能：

批处理：将多个独立的加载请求合并为单个批量请求
缓存：在单个请求生命周期内缓存已加载的数据

适用场景

DataLoader 特别适合以下场景：

需要从数据库批量获取数据
需要调用第三方 API 获取数据
存在嵌套查询可能导致 N+1 问题的场景

基本使用方法

1. 定义数据模型

首先定义一个简单的用户类型：

import strawberry

@strawberry.type
class User:
    id: strawberry.ID

2. 创建批量加载函数

这个函数接收一组键值，返回对应的用户列表：

from typing import List

async def load_users(keys: List[int]) -> List[User]:
    # 实际应用中这里会查询数据库或调用API
    return [User(id=key) for key in keys]

3. 创建 DataLoader 实例

from strawberry.dataloader import DataLoader

loader = DataLoader(load_fn=load_users)

4. 使用 DataLoader 加载数据

# 加载单个用户
user = await loader.load(1)

# 批量加载多个用户
[user_a, user_b] = await asyncio.gather(loader.load(1), loader.load(2))

高级功能

错误处理

当某些键值无法找到对应数据时，可以返回错误：

async def load_users(keys: List[int]) -> List[Union[User, ValueError]]:
    def lookup(key: int) -> Union[User, ValueError]:
        if user := users_database.get(key):
            return user
        return ValueError("用户未找到")
    return [lookup(key) for key in keys]

自定义缓存键

默认使用输入值作为缓存键，但可以自定义：

def custom_cache_key(key):
    return key.id  # 使用对象的id属性作为缓存键

loader = DataLoader(load_fn=loader_fn, cache_key_fn=custom_cache_key)

缓存管理

# 清除单个键的缓存
loader.clear(1)

# 清除多个键的缓存
loader.clear_many([1, 2, 3])

# 清除所有缓存
loader.clear_all()

预加载数据

可以将外部查询的数据导入缓存：

# 从数据库获取所有用户
people = await database.get_all_people()

# 将数据预加载到DataLoader缓存
loader.prime_many({person.id: person for person in people})

自定义缓存实现

可以替换默认的内存缓存：

class RedisCache(AbstractCache):
    def __init__(self, redis_client):
        self.redis = redis_client
    
    def get(self, key):
        return self.redis.get(key)
    
    # 实现其他必要方法...

loader = DataLoader(load_fn=load_users, cache_map=RedisCache(redis_client))

与 GraphQL 集成的最佳实践

在上下文中创建 DataLoader

为避免长期缓存导致的数据不一致，应在请求上下文中创建 DataLoader：

class MyGraphQL(GraphQL):
    async def get_context(self, request, response):
        return {
            "user_loader": DataLoader(load_fn=load_users)
        }

@strawberry.type
class Query:
    @strawberry.field
    async def get_user(self, info: strawberry.Info, id: strawberry.ID) -> User:
        return await info.context["user_loader"].load(id)

处理嵌套查询

@strawberry.type
class Person:
    id: strawberry.ID
    friends_ids: strawberry.Private[List[strawberry.ID]]

    @strawberry.field
    async def friends(self) -> List[Person]:
        return await loader.load_many(self.friends_ids)