graphql-java で雑に pagination 実装できるようにした

yoshiori.hatenablog.com

の続きです。 GraphQL の pagination っていうか Connections の仕様はGraphQL Cursor Connections Specification で、まぁ、実装した。基本的にはDBへクエリ投げやすくするように first/after と last/before を offset/limit に変換して渡すようにしてます。

コメントにも書いたけど変換はこんな感じになる。一応エッジケース全部テストしたはず。

   * for Example
   * [0,1,2,3,4]
   *
   * ## first/after
   * first:2, after:1 -> [2,3] -> offset:2, limit:2,
   * first:3, after:0 -> [1,2,3] -> offset:1, limit:3
   * first:null, after:2 -> [3,4] -> offset:3, limit:DEFAULT_SIZE(20)
   *
   * ## last/before
   * last:2, before:3 -> [1,2] -> offset: 1, limit:2
   * last: 2, before: null -> [3,4] -> offset: 3, limit: 2
   * last: null, before: 3 -> [0,1,2] -> offset: 0, limit: 3
   * last 10, before: 2 -> [0,1] -> offset: 0, limit: 2

DataFetcher Interface を実装してるので多分各種WebFrameworkでも使えるはず（俺は DGS 使ってるのでそれだけ動作確認済み）

使い方はトータル件数と offset/limit 引数に処理する関数を渡すだけ

new OffsetLimitConnection<>(fooService.getTotal(), fooService::getFoos).get(env));

標準でこのくらい用意してあるんじゃないの？とか誰か参考実装とか公開してるんじゃないの？？とか思ったけど見つけられなかったので同じような人がいたら助かるように貼っておきます。

import static graphql.Assert.assertNotNull;
import static java.nio.charset.StandardCharsets.UTF_8;

import com.google.common.base.Objects;
import com.google.common.io.BaseEncoding;
import graphql.relay.Connection;
import graphql.relay.ConnectionCursor;
import graphql.relay.DefaultConnection;
import graphql.relay.DefaultEdge;
import graphql.relay.DefaultPageInfo;
import graphql.relay.Edge;
import graphql.schema.DataFetcher;
import graphql.schema.DataFetchingEnvironment;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.function.BiFunction;

public class OffsetLimitConnection<T> implements DataFetcher<Connection<T>> {

  static final int DEFAULT_SIZE = 20;
  final int total;
  final BiFunction<Integer, Integer, List<T>> func;

  public OffsetLimitConnection(Integer total, BiFunction<Integer, Integer, List<T>> func) {
    this.total = assertNotNull(total, () -> "total cannot be null");
    this.func = assertNotNull(func, () -> "function cannot be null");
  }

  /**
   * Convert to offset/limit from last/before or last/before.
   *
   * <pre>
   * for Example
   * [0,1,2,3,4]
   *
   * ## first/after
   * first:2, after:1 -> [2,3] -> offset:2, limit:2,
   * first:3, after:0 -> [1,2,3] -> offset:1, limit:3
   * first:null, after:2 -> [3,4] -> offset:3, limit:DEFAULT_SIZE(20)
   *
   * ## last/before
   * last:2, before:3 -> [1,2] -> offset: 1, limit:2
   * last: 2, before: null -> [3,4] -> offset: 3, limit: 2
   * last: null, before: 3 -> [0,1,2] -> offset: 0, limit: 3
   * last 10, before: 2 -> [0,1] -> offset: 0, limit: 2
   * </pre>
   *
   * @param total
   * @param first
   * @param after
   * @param last
   * @param before
   * @return offsetAndLimit
   */
  static OffsetAndLimit getOffsetAndLimit(
      int total, Integer first, Integer after, Integer last, Integer before) {
    if ((first == null && after == null && last == null && before == null)
        || ((first != null || after != null) && (last != null || before != null))) {
      throw new IllegalArgumentException(
          "Only the combination of first/after or last/before can be given.");
    }
    int offset;
    int limit;

    // first/after
    if (first != null || after != null) {
      offset = after == null ? 0 : after + 1;
      limit = first == null ? DEFAULT_SIZE : first;

      // last/before
    } else {
      if (before == null) {
        before = total;
      }

      if (last == null) {
        last = DEFAULT_SIZE;
      }

      offset = Math.max(0, before - last);
      limit = Math.min(before, last);
    }
    return new OffsetAndLimit(offset, limit);
  }

  @Override
  public Connection<T> get(DataFetchingEnvironment environment) {
    Integer first = environment.getArgument("first");
    Integer after = NumberCursor.convertToOffset(environment.getArgument("after"));
    Integer last = environment.getArgument("last");
    Integer before = NumberCursor.convertToOffset(environment.getArgument("before"));

    var offsetAndLimit = getOffsetAndLimit(total, first, after, last, before);
    var offsetCursor = offsetAndLimit.offset;
    var limitCursor = offsetAndLimit.limit;

    List<T> list = func.apply(offsetCursor.toInt(), limitCursor.toInt());
    if (list.isEmpty()) {
      return emptyConnection();
    }

    List<Edge<T>> edges = new ArrayList<>();
    int index = offsetCursor.toInt();
    for (T o : list) {
      edges.add(new DefaultEdge<>(o, new NumberCursor(index++)));
    }

    var firstEdge = edges.get(0);
    var lastEdge = edges.get(edges.size() - 1);
    var pageInfo =
        new DefaultPageInfo(
            firstEdge.getCursor(),
            lastEdge.getCursor(),
            0 < ((NumberCursor) firstEdge.getCursor()).toInt(),
            ((NumberCursor) lastEdge.getCursor()).toInt() < total - 1);
    return new DefaultConnection<>(edges, pageInfo);
  }

  Connection<T> emptyConnection() {
    return new DefaultConnection<>(
        Collections.emptyList(), new DefaultPageInfo(null, null, false, false));
  }

  static class OffsetAndLimit {
    final NumberCursor offset;
    final NumberCursor limit;

    OffsetAndLimit(int offset, int limit) {
      this.offset = new NumberCursor(offset);
      this.limit = new NumberCursor(limit);
    }
  }

  static class NumberCursor implements ConnectionCursor {
    static final String PREFIX = "DUMMY";
    final int num;

    NumberCursor(int num) {
      this.num = num;
    }

    static Integer convertToOffset(String cursorString) {
      if (cursorString == null) {
        return null;
      }
      try {
        var string = new String(BaseEncoding.base64().decode(cursorString), UTF_8);
        return Integer.parseInt(string.substring(PREFIX.length()));
      } catch (IllegalArgumentException ignored) {
      }
      return null;
    }

    /** @return an opaque string that represents this cursor. */
    @Override
    public String getValue() {
      return convertToCursorString(num);
    }

    private String convertToCursorString(int num) {
      return BaseEncoding.base64().encode((PREFIX + num).getBytes(UTF_8));
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }
      NumberCursor that = (NumberCursor) o;
      return num == that.num;
    }

    @Override
    public int hashCode() {
      return Objects.hashCode(num);
    }

    public int toInt() {
      return num;
    }
    /**
     * In fact, the getValue method is not called, because toString has been called.
     *
     * @see <a href="https://github.com/graphql-java/graphql-java/issues/1944">GraphQL-java uses
     *     toString instead of ConnectionCursor getValue</a>
     * @return getValue() string
     */
    @Override
    public String toString() {
      return getValue();
    }
  }
}