GraphQL优化数据获取：graphql-java中Resolver合并子查询实现

给定以下GraphQL Schema：

type Node { 
  index: Int! 
  color: String! 
  neighbors(first: Int = null): [Node!]! 
} 
type Query { 
  nodes(color: String!): [Node!]! 
} 
schema { 
  query: Query 
}

需要执行的查询：

{ 
  nodes(color: "red") { 
    index 
    neighbors(first: 5) { 
      index 
    } 
  }
}

节点和邻居关系存储在不同的数据库表中，需要编写Resolver以最优方式获取数据，避免N+1查询，实现每个Resolver生成数据库子查询并有效合并，一次性获取数据。

这是典型的GraphQL N+1查询问题，在graphql-java里我们可以通过DataLoader批量加载机制解决，核心思路是先收集所有需要关联查询的父节点ID，然后一次性执行子查询获取所有关联数据，再映射回对应的父节点。下面是完整的实现步骤和代码示例：

1. 定义数据实体 #

先创建对应数据库表的实体类：

// 节点表实体
public class Node {
    private final Integer index;
    private final String color;

    public Node(Integer index, String color) {
        this.index = index;
        this.color = color;
    }

    // Getter方法
    public Integer getIndex() { return index; }
    public String getColor() { return color; }
}

// 邻居关联表实体（存储节点间的关联关系）
public class NeighborRelation {
    private final Integer sourceIndex;
    private final Integer targetIndex;

    public NeighborRelation(Integer sourceIndex, Integer targetIndex) {
        this.sourceIndex = sourceIndex;
        this.targetIndex = targetIndex;
    }

    // Getter方法
    public Integer getSourceIndex() { return sourceIndex; }
    public Integer getTargetIndex() { return targetIndex; }
}

2. 实现批量加载器（BatchLoader） #

创建一个BatchLoader来处理邻居数据的批量查询，它会接收所有需要查询邻居的源节点索引，一次性从数据库获取所有关联关系和目标节点数据：

import org.dataloader.BatchLoader;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.stream.Collectors;

public class NeighborsBatchLoader implements BatchLoader<Integer, List<Node>> {

    private final NodeRepository nodeRepository;
    private final Executor executor;

    public NeighborsBatchLoader(NodeRepository nodeRepository, Executor executor) {
        this.nodeRepository = nodeRepository;
        this.executor = executor;
    }

    @Override
    public CompletableFuture<List<List<Node>>> load(List<Integer> sourceNodeIndexes) {
        // 步骤1：批量查询所有源节点的邻居关联关系
        List<NeighborRelation> allRelations = nodeRepository.findNeighborRelationsBySourceIndexes(sourceNodeIndexes);
        
        // 步骤2：提取所有目标节点的索引，批量查询节点数据
        List<Integer> targetIndexes = allRelations.stream()
                .map(NeighborRelation::getTargetIndex)
                .distinct()
                .collect(Collectors.toList());
        Map<Integer, Node> targetNodeMap = nodeRepository.findNodesByIndexes(targetIndexes)
                .stream()
                .collect(Collectors.toMap(Node::getIndex, node -> node));
        
        // 步骤3：将源节点索引映射到对应的邻居列表，处理分页参数
        return CompletableFuture.supplyAsync(() -> 
            sourceNodeIndexes.stream()
                    .map(sourceIndex -> 
                        allRelations.stream()
                                .filter(rel -> rel.getSourceIndex().equals(sourceIndex))
                                .map(rel -> targetNodeMap.get(rel.getTargetIndex()))
                                .limit(5) // 先默认处理first=5，后续可优化为动态参数
                                .collect(Collectors.toList())
                    )
                    .collect(Collectors.toList()),
            executor
        );
    }
}

3. 编写Resolver实现 #

3.1 Query Resolver（根查询） #

负责处理nodes(color: String!)查询，一次性获取所有符合颜色条件的节点：

import graphql.kickstart.tools.GraphQLQueryResolver;
import org.springframework.stereotype.Component;
import java.util.List;

@Component
public class QueryResolver implements GraphQLQueryResolver {

    private final NodeRepository nodeRepository;

    public QueryResolver(NodeRepository nodeRepository) {
        this.nodeRepository = nodeRepository;
    }

    public List<Node> nodes(String color) {
        // 一次性查询所有红色节点，无N+1问题
        return nodeRepository.findNodesByColor(color);
    }
}

3.2 Node Resolver（字段Resolver） #

负责处理Node类型的neighbors字段，使用DataLoader批量加载邻居数据：

import graphql.kickstart.tools.GraphQLResolver;
import org.dataloader.DataLoader;
import org.springframework.stereotype.Component;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;

@Component
public class NodeResolver implements GraphQLResolver<Node> {

    private final DataLoader<Integer, List<Node>> neighborsDataLoader;

    public NodeResolver(DataLoader<Integer, List<Node>> neighborsDataLoader) {
        this.neighborsDataLoader = neighborsDataLoader;
    }

    public CompletableFuture<List<Node>> neighbors(Node node, Integer first) {
        // 使用DataLoader批量加载当前节点的邻居，而非单独查询
        return neighborsDataLoader.load(node.getIndex())
                .thenApply(neighbors -> neighbors.stream()
                        .limit(first != null ? first : 5)
                        .collect(Collectors.toList()));
    }
}

4. 配置DataLoader和GraphQL上下文 #

需要将DataLoader注册到DataLoaderRegistry，并确保每个请求上下文都能访问到它：

import org.dataloader.DataLoader;
import org.dataloader.DataLoaderRegistry;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import java.util.concurrent.Executors;

@Configuration
public class DataLoaderConfig {

    @Bean
    public DataLoaderRegistry dataLoaderRegistry(NodeRepository nodeRepository) {
        DataLoaderRegistry registry = new DataLoaderRegistry();
        // 创建并注册邻居数据加载器
        NeighborsBatchLoader batchLoader = new NeighborsBatchLoader(nodeRepository, Executors.newFixedThreadPool(4));
        DataLoader<Integer, List<Node>> neighborsDataLoader = DataLoader.newDataLoader(batchLoader);
        registry.register("neighborsDataLoader", neighborsDataLoader);
        return registry;
    }
}

然后配置GraphQL上下文构建器，让每个请求都能拿到DataLoader：

import graphql.kickstart.servlet.context.DefaultGraphQLServletContext;
import graphql.kickstart.servlet.context.GraphQLServletContextBuilder;
import org.dataloader.DataLoaderRegistry;
import org.springframework.stereotype.Component;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.websocket.Session;
import javax.websocket.server.HandshakeRequest;

@Component
public class CustomGraphQLContextBuilder implements GraphQLServletContextBuilder {

    private final DataLoaderRegistry dataLoaderRegistry;

    public CustomGraphQLContextBuilder(DataLoaderRegistry dataLoaderRegistry) {
        this.dataLoaderRegistry = dataLoaderRegistry;
    }

    @Override
    public DefaultGraphQLServletContext build(HttpServletRequest request, HttpServletResponse response) {
        return DefaultGraphQLServletContext.createServletContext()
                .with(dataLoaderRegistry)
                .build();
    }

    @Override
    public DefaultGraphQLServletContext build(Session session, HandshakeRequest request) {
        return DefaultGraphQLServletContext.createWebSocketContext()
                .with(dataLoaderRegistry)
                .build();
    }
}

5. 优化：动态处理first参数 #

上面的示例默认处理了first=5，如果需要动态支持传入的first参数，可以修改BatchLoader的Key为包含源节点索引和分页参数的自定义对象：

import java.util.Objects;

// 自定义Key类，包含源节点索引和分页参数
public class NeighborLoadKey {
    private final Integer sourceIndex;
    private final Integer first;

    public NeighborLoadKey(Integer sourceIndex, Integer first) {
        this.sourceIndex = sourceIndex;
        this.first = first != null ? first : 5;
    }

    // 必须正确实现equals和hashCode方法，确保DataLoader能正确缓存和分组
    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        NeighborLoadKey that = (NeighborLoadKey) o;
        return sourceIndex.equals(that.sourceIndex) && first.equals(that.first);
    }

    @Override
    public int hashCode() {
        return Objects.hash(sourceIndex, first);
    }

    // Getter方法
    public Integer getSourceIndex() { return sourceIndex; }
    public Integer getFirst() { return first; }
}

然后修改BatchLoader和Resolver：

// 修改后的BatchLoader
public class NeighborsBatchLoader implements BatchLoader<NeighborLoadKey, List<Node>> {

    @Override
    public CompletableFuture<List<List<Node>>> load(List<NeighborLoadKey> keys) {
        // 分组：按源节点索引和分页参数整理查询需求
        Map<Integer, Integer> sourceFirstMap = keys.stream()
                .collect(Collectors.toMap(NeighborLoadKey::getSourceIndex, NeighborLoadKey::getFirst));
        List<Integer> sourceIndexes = new ArrayList<>(sourceFirstMap.keySet());

        // 查询所有关联关系并按源节点分组
        Map<Integer, List<NeighborRelation>> sourceRelationsMap = nodeRepository.findNeighborRelationsBySourceIndexes(sourceIndexes)
                .stream()
                .collect(Collectors.groupingBy(NeighborRelation::getSourceIndex));

        // 批量查询目标节点
        List<Integer> targetIndexes = sourceRelationsMap.values().stream()
                .flatMap(List::stream)
                .map(NeighborRelation::getTargetIndex)
                .distinct()
                .collect(Collectors.toList());
        Map<Integer, Node> targetNodeMap = nodeRepository.findNodesByIndexes(targetIndexes)
                .stream()
                .collect(Collectors.toMap(Node::getIndex, node -> node));

        // 映射每个Key到对应的分页后的邻居列表
        return CompletableFuture.supplyAsync(() -> 
            keys.stream()
                    .map(key -> {
                        List<NeighborRelation> relations = sourceRelationsMap.getOrDefault(key.getSourceIndex(), List.of());
                        return relations.stream()
                                .map(rel -> targetNodeMap.get(rel.getTargetIndex()))
                                .limit(key.getFirst())
                                .collect(Collectors.toList());
                    })
                    .collect(Collectors.toList()),
            executor
        );
    }
}

// 修改后的NodeResolver
public CompletableFuture<List<Node>> neighbors(Node node, Integer first) {
    NeighborLoadKey key = new NeighborLoadKey(node.getIndex(), first);
    return neighborsDataLoader.load(key);
}

工作原理说明 #

1. 根查询nodes(color: "red")会一次性从数据库获取所有红色节点，这是第1次查询。
2. GraphQL引擎会遍历所有红色节点，收集它们的neighbors字段查询请求，DataLoader会将这些请求合并成一个批量查询。
3. BatchLoader执行第2次查询：一次性获取所有红色节点的邻居关联关系，再批量获取对应的目标节点数据，最后映射回每个源节点。
4. 整个过程仅需2次数据库查询，彻底避免了N+1问题，性能最优。

内容的提问来源于stack exchange，提问作者Suhas Hebbar