DGS Framework：如何向数据获取器传递额外数据？

我有如下@DgsQuery方法：

@DgsQuery
public List<Sensor> sensors(DgsDataFetchingEnvironment dfe) {
     
    SensorType sensorType = SensorType.newBuilder()
            .id(UUID.randomUUID().toString())
            .build();
    

    Sensor sensor = new Sensor();
    sensor.setId(UUID.randomUUID().toString());
    sensor.setName("foo");
    sensor.setSensorType(sensorType);

    return Arrays.asList(sensor);
}

以及对应的@DgsData数据获取器：

@DgsData(parentType = DgsConstants.SENSOR.TYPE_NAME, field = DgsConstants.SENSOR.SensorType)
public CompletableFuture<SensorType> sensorType(DgsDataFetchingEnvironment dfe) {
    Sensor sensor = dfe.getSource();
    UUID sensorTypeId = UUID.fromString(sensor.getSensorType().getId());

    // fetch sensorType by ID from DB, map to generated class and return
    LOGGER.info("Would now fetch sensorType with ID {} via RPC or whatever you want ;)", sensorTypeId);

    return null;
}

现在需要在数据获取器中访问@DgsQuery方法里初始化的额外数据/对象（比如从数据库加载的Entity），但目前只能通过dfe.getSource()拿到Sensor对象，无法获取其他已加载的额外数据，该如何实现？

1. 扩展DTO对象，添加内部字段存储额外数据 #

如果允许修改Sensor这类DTO，可以添加一个仅用于内部传递的字段（无需在GraphQL Schema中暴露），直接存储额外数据。

修改@DgsQuery方法：

@DgsQuery
public List<Sensor> sensors(DgsDataFetchingEnvironment dfe) {
    // 从数据库加载目标Entity
    SensorEntity sensorEntity = sensorRepository.findById(someId);

    SensorType sensorType = SensorType.newBuilder()
            .id(sensorEntity.getSensorTypeId().toString())
            .build();
    
    Sensor sensor = new Sensor();
    sensor.setId(sensorEntity.getId().toString());
    sensor.setName(sensorEntity.getName());
    sensor.setSensorType(sensorType);
    // 将额外Entity存入DTO的内部字段
    sensor.setSensorEntity(sensorEntity);

    return Arrays.asList(sensor);
}

在数据获取器中直接读取：

@DgsData(parentType = DgsConstants.SENSOR.TYPE_NAME, field = DgsConstants.SENSOR.SensorType)
public CompletableFuture<SensorType> sensorType(DgsDataFetchingEnvironment dfe) {
    Sensor sensor = dfe.getSource();
    // 获取额外存储的Entity数据
    SensorEntity sensorEntity = sensor.getSensorEntity();
    
    // 直接使用Entity数据构建结果，无需重复查库
    SensorType sensorType = SensorType.newBuilder()
            .id(sensorEntity.getSensorTypeId().toString())
            .name(sensorEntity.getSensorTypeName())
            .build();

    return CompletableFuture.completedFuture(sensorType);
}

2. 利用DgsDataFetchingEnvironment上下文存储 #

通过DgsDataFetchingEnvironment的setContext()和getContext()方法，在Query阶段将额外数据存入上下文，后续数据获取器从上下文读取。

方式1：用Map存储多数据 #

修改@DgsQuery方法：

@DgsQuery
public List<Sensor> sensors(DgsDataFetchingEnvironment dfe) {
    SensorEntity sensorEntity = sensorRepository.findById(someId);

    // 用Map存储额外数据，以Sensor ID作为键区分多对象
    Map<String, Object> contextMap = new HashMap<>();
    contextMap.put("sensorEntity_" + sensorEntity.getId(), sensorEntity);
    dfe.setContext(contextMap);

    SensorType sensorType = SensorType.newBuilder()
            .id(sensorEntity.getSensorTypeId().toString())
            .build();
    
    Sensor sensor = new Sensor();
    sensor.setId(sensorEntity.getId().toString());
    sensor.setName(sensorEntity.getName());
    sensor.setSensorType(sensorType);

    return Arrays.asList(sensor);
}

数据获取器中读取：

@DgsData(parentType = DgsConstants.SENSOR.TYPE_NAME, field = DgsConstants.SENSOR.SensorType)
public CompletableFuture<SensorType> sensorType(DgsDataFetchingEnvironment dfe) {
    Sensor sensor = dfe.getSource();
    Map<String, Object> contextMap = (Map<String, Object>) dfe.getContext();
    // 根据Sensor ID匹配对应的Entity
    SensorEntity sensorEntity = (SensorEntity) contextMap.get("sensorEntity_" + sensor.getId());
    
    SensorType sensorType = SensorType.newBuilder()
            .id(sensorEntity.getSensorTypeId().toString())
            .name(sensorEntity.getSensorTypeName())
            .build();

    return CompletableFuture.completedFuture(sensorType);
}

方式2：自定义上下文类（类型更安全） #

定义专属上下文类：

public class SensorQueryContext {
    private Map<UUID, SensorEntity> sensorEntityMap;

    // 构造器、getter、setter方法
}

修改@DgsQuery方法：

@DgsQuery
public List<Sensor> sensors(DgsDataFetchingEnvironment dfe) {
    SensorEntity sensorEntity = sensorRepository.findById(someId);

    SensorQueryContext context = new SensorQueryContext();
    context.setSensorEntityMap(Collections.singletonMap(sensorEntity.getId(), sensorEntity));
    dfe.setContext(context);

    // 剩余Sensor、SensorType初始化代码不变
    // ...

    return Arrays.asList(sensor);
}

数据获取器中读取：

@DgsData(parentType = DgsConstants.SENSOR.TYPE_NAME, field = DgsConstants.SENSOR.SensorType)
public CompletableFuture<SensorType> sensorType(DgsDataFetchingEnvironment dfe) {
    Sensor sensor = dfe.getSource();
    SensorQueryContext context = (SensorQueryContext) dfe.getContext();
    SensorEntity sensorEntity = context.getSensorEntityMap().get(UUID.fromString(sensor.getId()));
    
    // 构建并返回SensorType
    // ...
}

3. 用DataLoader预加载缓存（批量场景推荐） #

如果是批量数据场景，推荐使用DGS的DataLoader机制，既可以避免N+1查询，也能复用Query阶段已加载的数据。

首先定义DataLoader：

@DgsDataLoader(name = "sensorTypeLoader")
public class SensorTypeDataLoader extends BatchLoader<UUID, SensorType> {
    private final SensorRepository sensorRepository;

    public SensorTypeDataLoader(SensorRepository sensorRepository) {
        this.sensorRepository = sensorRepository;
    }

    @Override
    public CompletionStage<List<SensorType>> load(List<UUID> sensorTypeIds) {
        // 批量查询数据并转换为SensorType
        List<SensorEntity> entities = sensorRepository.findBySensorTypeIdIn(sensorTypeIds);
        List<SensorType> sensorTypes = entities.stream()
                .map(e -> SensorType.newBuilder()
                        .id(e.getSensorTypeId().toString())
                        .name(e.getSensorTypeName())
                        .build())
                .collect(Collectors.toList());
        return CompletableFuture.completedFuture(sensorTypes);
    }
}

在@DgsQuery中预加载数据到DataLoader缓存：

@DgsQuery
public List<Sensor> sensors(DgsDataFetchingEnvironment dfe) {
    SensorEntity sensorEntity = sensorRepository.findById(someId);

    // 获取DataLoader并将已加载的数据存入缓存
    DataLoader<UUID, SensorType> loader = dfe.getDataLoader("sensorTypeLoader");
    loader.prime(sensorEntity.getSensorTypeId(), 
            SensorType.newBuilder()
                    .id(sensorEntity.getSensorTypeId().toString())
                    .name(sensorEntity.getSensorTypeName())
                    .build());

    // 剩余Sensor、SensorType初始化代码不变
    // ...

    return Arrays.asList(sensor);
}

数据获取器中直接使用DataLoader（自动命中缓存）：

@DgsData(parentType = DgsConstants.SENSOR.TYPE_NAME, field = DgsConstants.SENSOR.SensorType)
public CompletableFuture<SensorType> sensorType(DgsDataFetchingEnvironment dfe) {
    Sensor sensor = dfe.getSource();
    UUID sensorTypeId = UUID.fromString(sensor.getSensorType().getId());
    
    // 直接从DataLoader获取，优先使用预加载的缓存数据
    return dfe.getDataLoader("sensorTypeLoader").load(sensorTypeId);
}

内容的提问来源于stack exchange，提问作者Alex Tbk