Introduction

While GraphQL federation enables team to collaborate on different subgraphs, a problem arises: how does a subgraph handle a N+1 queries problem? For example, in our previous post, if we have a list of products, each of which has multiple reviews, then it can lead to multiple queries in the Review subgraph as it needs to resolve reviews field for each Product entity one by one. In this post, we’ll explore a pattern called Data Loader, that’s often used in a federated GraphQL architecture to solve the N+1 queries problem.

We’ll only go through the important bits, but the final runnable project can be browsed here.

Setup

To illustrate the problems better, we’ll change the setup by introducing a MySQL database to host our reviews data instead of keeping them in memory like in the original post.

The first thing to do is to implement a MySQL repository that allows us to query reviews by product ID.

type Repository struct {
	db *sql.DB
}

func NewRepository(db *sql.DB) *Repository {
	return &Repository{db: db}
}

func (r *Repository) GetByProductID(ctx context.Context, productID string) ([]Review, error) {
	query := `SELECT id, body, product_id FROM reviews WHERE product_id = ?`

	rows, err := r.db.QueryContext(ctx, query, productID)
	if err != nil {
		return nil, fmt.Errorf("query failed: %w", err)
	}
	defer rows.Close()

	var reviews []Review
	for rows.Next() {
		var r Review
		if err := rows.Scan(&r.ID, &r.Body, &r.ProductID); err != nil {
			return nil, fmt.Errorf("row scan failed: %w", err)
		}
		reviews = append(reviews, r)
	}

	return reviews, nil
}

type Review struct {
	ID        string
	Body      string
	ProductID string
}

And then we need to make change to the resolver, so that it can use this repository to query reviews from the MySQL database.

func (r *entityResolver) FindProductByID(ctx context.Context, id string) (*model.Product, error) {
	reviews, err := r.ReviewsGetter.GetByProductID(id)
	if err != nil {
		return nil, err
	}

	var reviewModels []*model.Review
	for _, review := range reviews {
		reviewModels = append(reviewModels, &model.Review{
			ID:   review.ID,
			Body: review.Body,
		})
	}

    return &model.Product{
		ID:      id,
		Reviews: reviewModels,
	}, nil
}

So far so good, we can run make start to initialize all our containers, including all subgraphs and the Apollo router, and everything should work fine. However, it can be noticed that we’re doing multiple SELECT queries each time the router asks Review subgraph to get reviews of a product. When our throughput is large enough, this approach will become a bottle neck. Instead, it will be more efficient if we can group multiple product IDs from multiple entity resolver calls into one, and only execute one SELECT statement.

This introduces an implementation problem, because of the signature of the resolver.

func (r *entityResolver) FindProductByID(ctx context.Context, id string) (*model.Product, error)

This cannot be changed because each product entity must be resolved individually. Therefore, even if we introduce a new method in the repository that accepts a list of product ID, it may not be feasible here.

Data loader

Data loader is a pattern that’s often used in GraphQL Federation to solve this kind of problem. In a nutshell, a generic data loader groups different keys together and then perform an operation on these keys. The grouping policy can be based on batch size, e.g. grouping 10 keys together; or based on time, e.g. grouping keys within 10ms as a batch.

In reviews subgraph, we can implement a reviews data loader that groups multiple product ID together (each product ID can be considered a “key” here), and then only makes one database call to retrieve all reviews of those products.

For the purpose of this post, we will use the data loader implementation from github.com/vikstrous/dataloadgen. But first, we will still need a method in the reviews repository to fetch reviews for multiple products.

func (r *Repository) GetByProductIDs(ctx context.Context, productIDs []string) (map[string][]Review, error) {
	placeHolders := make([]string, 0, len(productIDs))
	args := make([]any, 0, len(productIDs))
	for _, productID := range productIDs {
		placeHolders = append(placeHolders, "?")
		args = append(args, productID)
	}

	query := fmt.Sprintf("SELECT id, body, product_id FROM reviews WHERE product_id IN (%s)", strings.Join(placeHolders, ","))

	rows, err := r.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("query failed: %w", err)
	}
	defer rows.Close()

	results := make(map[string][]Review)
	for rows.Next() {
		var r Review
		if err := rows.Scan(&r.ID, &r.Body, &r.ProductID); err != nil {
			return nil, fmt.Errorf("row scan failed: %w", err)
		}
		results[r.ProductID] = append(results[r.ProductID], r)
	}

	return results, nil
}

Then we will use the data loader to group multiple product ids before invoking the repository’s method.

func NewLoader(repo *storage.Repository) *dataloadgen.Loader[string, []storage.Review] {
	r := &reader{
		repo: repo,
	}
	return dataloadgen.NewLoader(r.getByProductIDs, dataloadgen.WithWait(time.Millisecond))
}

type reader struct {
	repo *storage.Repository
}

func (r *reader) getByProductIDs(ctx context.Context, productIDs []string) ([][]storage.Review, []error) {
	results, err := r.repo.GetByProductIDs(ctx, productIDs)
	if err != nil {
		return nil, []error{err}
	}

	reviews := make([][]storage.Review, 0)
	for _, productID := range productIDs {
		reviews = append(reviews, results[productID])
	}

	return reviews, nil
}

NewLoader creates a new instace of a reviews data loader. The only dependency it needs is an instance of the reviews repository. Under the hood, we’re calling dataloadgen.New to creates a new data loader instance, which is a generic function that accepts a fetch function and a list of optional function arguments. For the fetch function, we implement a reader type, which wraps around the reviews repository. reader provides a method getByProductIDs, which satisfies the fetch function signature requirement of dataloadgen.New while also calls method GetByProductIDs of reviews repository. When constructing a new dataloader, there are a few optional function arguments available to choose, but for this post, we’re only using dataloadgen.WithWait(time.Millisecond) which instructs the data loader to wait for 1 millisecond before invoking the fetch function.

How do we use this data loader in our GraphQL resolver? One approach is to create the data loader on a per request basis.

type ctxKey string

const loaderCtxKey ctxKey = "loader"

func HTTPMiddleware(repo *storage.Repository, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		l := NewLoader(repo)
		r = r.WithContext(context.WithValue(r.Context(), loaderCtxKey, l))
		next.ServeHTTP(w, r)
	})
}

func From(ctx context.Context) *dataloadgen.Loader[string, []storage.Review] {
	return ctx.Value(loaderCtxKey).(*dataloadgen.Loader[string, []storage.Review])
}

HTTPMiddleware returns a new HTTP handler that creates a new data loader instance being attached to each request’s context, which then can be obtained by calling the From method. In our resolver, we have to make some changes to leverage on these new utilities.

func (r *entityResolver) FindProductByID(ctx context.Context, id string) (*model.Product, error) {
	reviewsLoader := loader.From(ctx)
	reviews, err := reviewsLoader.Load(ctx, id)
	if err != nil {
		return nil, err
	}

	var reviewModels []*model.Review
	for _, review := range reviews {
		reviewModels = append(reviewModels, &model.Review{
			ID:   review.ID,
			Body: review.Body,
		})
	}

	return &model.Product{
		ID:      id,
		Reviews: reviewModels,
	}, nil
}

Instead of calling the reviews repository directly, we’re getting the reviews data loader first from context, and then call its Load method to fetch the reviews for the product ID we’re resolving. That’s the only change we need in our resolver.

Another approach is to have a single reviews data loader on the application level. This can be very powerful as it eliminates the need to have multiple instances for multiple requests, enables the ability to group product IDs of multiple requests into 1 query (further reducing the load for our database), and if it permits, the data loader can be set up to hold reviews data in memory, making it a pseudo cache layer. The tradeoff is that it makes operating the application more complicated: much attention needs to be paid on configurations such as cache TTL and how to manually cleaning up cache by calling the Clear method of data loader.

Summary

There may be a lot to take in in this post, and it takes a bit of a learning curve to get used to how data loader works. But as mentioned previously, the fully runnable demo project can be found here. It is possible to make some configuration changes for the data loader, like changing the number of max items per batch, and observe how the behavior of the application changes.