Spring GraphQL 1.0.0-M3でDataLoaderもアノテーションマッピングできるようになった

Spring GraphQL 1.0.0-M2がリリースされました。

• Spring GraphQL 1.0.0-M3 Released

N + 1問題へ対応できるDataLoaderもコントローラーへアノテーションマッピングができるようになったのでメモしておきたいと思います。

1.0.0-M2がリリースされたときに書いた記事では次のようなGraphQLスキーマとコントローラーを紹介しました。

type Query {
  comics(genre: String!): [Comic!]!
}

type Comic {
  id: ID!
  title: String!
  author: Author!
}

type Author {
  id: ID!
  name: String!
}

package com.example;

import java.util.List;
import java.util.Optional;

import org.springframework.graphql.data.method.annotation.Argument;
import org.springframework.graphql.data.method.annotation.QueryMapping;
import org.springframework.graphql.data.method.annotation.SchemaMapping;
import org.springframework.stereotype.Controller;

@Controller
public class ComicController {

    private final ComicRepository comicRepository;
    private final AuthorRepository authorRepository;

    public ComicController(ComicRepository comicRepository, AuthorRepository authorRepository) {
        this.comicRepository = comicRepository;
        this.authorRepository = authorRepository;
    }

    @QueryMapping
    public List<Comic> comics(@Argument String genre) {
        return comicRepository.findByGenre(genre);
    }

    @SchemaMapping
    public Optional<Author> author(Comic source) {
        return authorRepository.findById(source.getAuthorId());
    }
}

この場合genreを条件にしてComicを複数取得して、各Comic毎にAuthorを検索するといった動作をします。
つまりComicで1回、Authorでn回の検索が行われます。

これを解消するのがDataLoaderですが、1.0.0-M3から次のように書けるようになりました。

    @BatchMapping
    public List<Author> author(List<Comic> sources) {
        List<Integer> ids = sources.stream().map(Comic::getAuthorId).toList();
        Map<Integer, Author> authors = authorRepository.findByIds(ids)
                .stream().collect(toMap(Author::getId, identity()));
        return ids.stream().map(authors::get).toList();
    }

引数にComicのリストを取って複数件まとめて検索しています（authorRepository.findByIds(ids)で発行されるクエリはselect * from author where id in (...)だと思ってください）。

取得したAuthorのリストをマップに変換し、idsをもとにして再度リストを構築して返していますが、これは戻り値となるリストの要素が引数のそれに対応しているためです。
例えば引数が[ドラゴンボール, Dr.スランプ, 2.5次元の誘惑]だった場合、戻り値は[鳥山明, 鳥山明, 橋本悠]にならないといけないということです。

ちなみに戻り値はマップでもOKです。
その場合、戻り値の型はMap<Comic, Author>になります。

最後にソースコードを軽く読んだのでメモっておきます。

• @BatchMappingを読み取っているのはorg.springframework.graphql.data.method.annotation.support.AnnotatedControllerConfigurer
• データフェッチの実装はorg.springframework.graphql.data.method.annotation.support.AnnotatedControllerConfigurer.BatchMappingDataFetcher
• DataLoaderを定義しているのはorg.springframework.graphql.data.method.annotation.support.AnnotatedControllerConfigurerのregisterBatchLoaderメソッド
• リクエスト毎に使うDataLoaderを登録しているのはorg.springframework.graphql.execution.ExecutionGraphQlServiceのinitExecutionInputメソッド