Using java-dataloader to fix N+1 queries in a GraphQL API - Jordy van Vorselen

Using java-dataloader to fix N+1 queries in a GraphQL API - Jordy van Vorselen

SELECT id, name FROM authors WHERE id = 1;

SELECT id, name FROM books WHERE author_id = 1;
/* Query above gives the following book ID's: [1, 2, 3] */

SELECT id, name FROM reviews WHERE book_id = 1;
SELECT id, name FROM reviews WHERE book_id = 2;
SELECT id, name FROM reviews WHERE book_id = 3;
/* etc... */

/* For each ID, a separate query is executed */

@Component
public class DataLoaderRegistryFactory {

  private final ReviewRepository reviewRepository;

  // Needed because postgres allows a maximum of ~30000 query parameters.
  private static final int MAX_BATCH_SIZE = 30000;

  public DataLoaderRegistryFactory(ReviewRepository reviewRepository) {
    this.reviewRepository = reviewRepository;
  }

  public DataLoaderRegistry build() {
    var options = DataLoaderOptions.newOptions().setMaxBatchSize(MAX_BATCH_SIZE);

    var reviewDataLoader = DataLoaderFactory.newMappedDataLoader(new ReviewsBatchLoader(reviewRepository), options);

    return DataLoaderRegistry.newRegistry()
            .register(ReviewBatchLoader.KEY, reviewDataLoader)
            .build();
  }
}

public class ReviewsBatchLoader implements MappedBatchLoader<Integer, List<Review>> {

  public static final String KEY = "REVIEWS_BATCH_LOADER";

  private final ReviewRepository reviewRepository;

  public ReviewsBatchLoader(ReviewRepository reviewRepository) {
    this.reviewRepository = reviewRepository;
  }

  @Override
  public CompletionStage<Map<Integer, List<Review>>> load(Set<Integer> bookIds) {
    return CompletableFuture.supplyAsync(() -> {
      // If possible, return a map of { id: [Review] } straight out of the repository to avoid
      // having to build the map yourself.
      var reviews = reviewRepository.findAllByBookIdIn(bookIds.stream().toList());

      // If it's not possible to return a map directly from the repository, you can use something
      // like this to build the map yourself.
      return bookIds.stream()
              .collect(
                      HashMap::new,
                      (map, bookId) -> map.put(
                              bookId,
                              reviews.stream()
                                      .filter(review -> Objects.equals(bookId, review.getBook().getId()))
                                      .toList()
                      ),
                      HashMap::putAll
              );
    });
  }
}

@Component
@RequiredArgsConstructor
public class BookResolver implements GraphQLResolver<Review> {
  public List<Review> getReviews(Book book, DataFetchingEnvironment environment) {
    var registry = environment.getDataLoaderRegistry();
    DataLoader<Integer, List<Review>> dataLoader = registry.getDataLoader(ReviewsBatchLoader.KEY);

    return dataLoader.load(book.getId());
  }
}