如何用Golang的gqlgen与MongoDB实现GraphQL Data Loader

针对你遇到的产品分类N+1查询问题，我们可以通过graph-gophers/dataloader库结合MongoDB的批量查询能力来解决，以下是具体实现步骤：

1. 定义分类的Data Loader #

首先创建Data Loader的批量加载逻辑，利用MongoDB的$in操作符一次性查询所有需要的分类：

package dataloaders

import (
	"context"
	"fmt"
	"github.com/graph-gophers/dataloader"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"time"
)

// Category 映射MongoDB分类文档结构
type Category struct {
	ID   primitive.ObjectID `bson:"_id"`
	Name string             `bson:"name"`
}

// CategoryKey 作为Data Loader的键类型，基于ObjectID
type CategoryKey primitive.ObjectID

func (k CategoryKey) String() string {
	return primitive.ObjectID(k).Hex()
}

// NewCategoryLoader 创建分类Data Loader实例
func NewCategoryLoader(col *mongo.Collection) *dataloader.Loader {
	batchLoadFn := func(ctx context.Context, keys dataloader.Keys) []*dataloader.Result {
		// 提取所有需要查询的分类ID
		var categoryIDs []primitive.ObjectID
		for _, key := range keys {
			id, ok := key.Raw().(primitive.ObjectID)
			if !ok {
				continue
			}
			categoryIDs = append(categoryIDs, id)
		}

		// 批量查询MongoDB分类
		var categories []Category
		cursor, err := col.Find(ctx, bson.M{"_id": bson.M{"$in": categoryIDs}})
		if err != nil {
			// 错误处理：给每个请求返回相同错误
			results := make([]*dataloader.Result, len(keys))
			for i := range results {
				results[i] = &dataloader.Result{Error: err}
			}
			return results
		}
		defer cursor.Close(ctx)

		if err := cursor.All(ctx, &categories); err != nil {
			results := make([]*dataloader.Result, len(keys))
			for i := range results {
				results[i] = &dataloader.Result{Error: err}
			}
			return results
		}

		// 构建ID到分类的映射表
		categoryMap := make(map[primitive.ObjectID]Category)
		for _, cat := range categories {
			categoryMap[cat.ID] = cat
		}

		// 按请求的ID顺序返回结果
		results := make([]*dataloader.Result, len(keys))
		for i, key := range keys {
			id, ok := key.Raw().(primitive.ObjectID)
			if !ok {
				results[i] = &dataloader.Result{Error: fmt.Errorf("无效的分类ID类型")}
				continue
			}
			cat, exists := categoryMap[id]
			if !exists {
				results[i] = &dataloader.Result{Error: fmt.Errorf("未找到分类: %s", id.Hex())}
				continue
			}
			results[i] = &dataloader.Result{Data: cat}
		}
		return results
	}

	// 配置批量加载的参数：批量容量、等待时间（聚合请求的窗口）
	return dataloader.NewBatchedLoader(
		batchLoadFn,
		dataloader.WithBatchCapacity(100),
		dataloader.WithWait(time.Millisecond*10),
	)
}

2. 将Data Loader注入GraphQL上下文 #

在gqlgen的服务初始化中，把Data Loader实例注入到请求上下文，确保每个请求都有独立的Loader实例：

package server

import (
	"context"
	"net/http"
	"your-project/dataloaders"
	"go.mongodb.org/mongo-driver/mongo"
)

// 自定义上下文键，避免冲突
type ctxKey string

const categoryLoaderKey ctxKey = "categoryLoader"

// ContextMiddleware 中间件：将分类Data Loader注入上下文
func ContextMiddleware(categoryCol *mongo.Collection) func(next http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// 为当前请求创建Data Loader
			loader := dataloaders.NewCategoryLoader(categoryCol)
			// 注入上下文
			ctx := context.WithValue(r.Context(), categoryLoaderKey, loader)
			// 传递上下文到后续处理
			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}

// GetCategoryLoader 从上下文获取分类Data Loader
func GetCategoryLoader(ctx context.Context) *dataloader.Loader {
	loader, ok := ctx.Value(categoryLoaderKey).(*dataloader.Loader)
	if !ok {
		panic("上下文未找到分类Data Loader")
	}
	return loader
}

3. 在Resolver中使用Data Loader #

修改产品的Category resolver，通过Data Loader获取分类，替代单独查询：

package resolver

import (
	"context"
	"your-project/gqlgen"
	"your-project/server"
	"your-project/dataloaders"
	"go.mongodb.org/mongo-driver/bson/primitive"
)

// ProductResolver 产品Resolver实现
type ProductResolver struct{}

// Category 获取产品关联的分类
func (r *ProductResolver) Category(ctx context.Context, obj *gqlgen.Product) (*gqlgen.Category, error) {
	// 转换产品中的分类ID为ObjectID
	catID, err := primitive.ObjectIDFromHex(obj.Category)
	if err != nil {
		return nil, err
	}

	// 从上下文获取Data Loader
	loader := server.GetCategoryLoader(ctx)
	// 加载分类（自动批量处理）
	result, err := loader.Load(ctx, dataloaders.CategoryKey(catID))()
	if err != nil {
		return nil, err
	}

	// 转换为GraphQL定义的Category类型
	cat := result.(dataloaders.Category)
	return &gqlgen.Category{
		ID:   cat.ID.Hex(),
		Name: cat.Name,
	}, nil
}

工作原理 #

• Data Loader会在短时间窗口（比如10ms）内收集所有分类查询请求，把重复或批量的ID合并成一次MongoDB查询（用$in）。
• 同一个请求内的重复分类ID会直接返回缓存结果，无需重复查询。
• 原本1000次的分类查询会被压缩成1次批量查询，彻底解决N+1问题。

内容的提问来源于stack exchange，提问作者Ketan Rathod