JPAQueryGraphQLConnection.java

import com.querydsl.jpa.impl.JPAQuery;
import graphql.schema.DataFetcher;
import graphql.schema.DataFetchingEnvironment;
import graphql.schema.GraphQLList;
import graphql.schema.GraphQLObjectType;
import lombok.Builder;
import lombok.RequiredArgsConstructor;
import lombok.Value;

import java.util.Base64;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static graphql.Scalars.*;
import static graphql.schema.GraphQLFieldDefinition.newFieldDefinition;
import static graphql.schema.GraphQLObjectType.newObject;

/**
 * Created by joeyang on 27/02/2017.
 */
public class JPAQueryGraphQLConnection<T> implements DataFetcher {

    JPAQuery<T> jpaQuery;

    public JPAQueryGraphQLConnection(JPAQuery<T> jpaQuery) {
        this.jpaQuery = jpaQuery;
    }

    @Value
    public static class Cursor {

        public static Cursor of(long offset) {
            String value = offsetCursorConverter.fromOffset(offset);
            return new Cursor(value);
        }

        String value;

        @Override
        public String toString() {
            return value;
        }
    }

    @Value
    @RequiredArgsConstructor
    public static class PageInfo {
        boolean hasNextPage;
        Cursor endCursor;
    }

    @Value
    @RequiredArgsConstructor
    public static class Edge<T> {
        T node;
        Cursor cursor;
    }

    @Value
    @RequiredArgsConstructor
    @Builder
    public static class Connection<T> {
        long totalCount;
        List<Edge<T>> edges;
        PageInfo pageInfo;
    }


    private static GraphQLObjectType edgeType(GraphQLObjectType nodeType) {

        return newObject()
            .name(nodeType.getName() + "Edge")
            .field(newFieldDefinition()
                .type(nodeType)
                .name("node"))
            .field(newFieldDefinition()
                .type(GraphQLString)
                .name("cursor"))
            .build();

    }

    public static GraphQLObjectType connectionType(GraphQLObjectType nodeType) {

        GraphQLObjectType pageType = newObject()
            .name("pageInfo")
            .field(newFieldDefinition()
                .name("endCursor")
                .type(GraphQLString))
            .field(newFieldDefinition()
                .name("hasNextPage")
                .type(GraphQLBoolean))
            .build();

        return newObject()
            .name(nodeType.getName() + "Connection")
            .field(newFieldDefinition()
                .type(GraphQLLong)
                .name("totalCount"))
            .field(newFieldDefinition()
                .type(new GraphQLList(edgeType(nodeType)))
                .name("edges"))
            .field(newFieldDefinition()
                .name("pageInfo")
                .type(pageType))
            .build();
    }

    public static class OffsetCursorConverter {

        private static final String PREFIX = "jpa-cursor";

        String fromOffset(long offset) {
            byte[] payload = (PREFIX + String.valueOf(offset)).getBytes();
            return Base64.getEncoder().encodeToString(payload);
        }

        int toOffset(String cursor) {
            String decoded = new String(Base64.getDecoder().decode(cursor));
            String value = decoded.substring(PREFIX.length());
            return Integer.parseInt(value);
        }

    }

    private static final OffsetCursorConverter offsetCursorConverter = new OffsetCursorConverter();

    @Override
    public Object get(DataFetchingEnvironment environment) {

        Integer limit = environment.getArgument("first");
        String after = environment.getArgument("after");

        int offset = after != null ? offsetCursorConverter.toOffset(after) + 1 : 0;

        JPAQuery<T> query = jpaQuery
            .offset(offset)
            .limit(limit);

        long totalCount = jpaQuery.fetchCount();
        List<T> results = query.fetch();

        boolean hasNextPage = (offset + limit) < totalCount;

        PageInfo pageInfo = new PageInfo(hasNextPage, Cursor.of(totalCount - 1));

        List<Edge<T>> edges = IntStream.range(0, results.size())
            .mapToObj(index -> new Edge<>(results.get(index), Cursor.of((long) index)))
            .collect(Collectors.toList());

        Connection<T> result = Connection.<T>builder()
            .totalCount(totalCount)
            .pageInfo(pageInfo)
            .edges(edges)
            .build();

        return result;
    }

}