How to respond with extensions using graphql-java?

The best way I’ve found to return extensions is to implement a subclass of SimpleInstrumentation that overrides instrumentExecutionResult (code stolen partially from graphql-java’s TracingInstrumentation):

@Override
public CompletableFuture<ExecutionResult> instrumentExecutionResult(
        ExecutionResult executionResult,
        InstrumentationExecutionParameters parameters) {
    Map<Object, Object> currentExt = executionResult.getExtensions();
    Map<Object, Object> newExtensionMap = new LinkedHashMap<>();
    newExtensionMap.putAll(currentExt == null ? Collections.emptyMap() : currentExt);
    newExtensionMap.put("MyExtensionKey", myExtensionValue);

    return CompletableFuture.completedFuture(
        new ExecutionResultImpl(
            executionResult.getData(), 
            executionResult.getErrors(), 
            newExtensionMap));
}

When setting up the GraphQL instance you then pass an instance of the instrumentation class in:

GraphQL graphQL = GraphQL
        .newGraphQL(schema)
        .instrumentation(new MyInstrumentation())
        .build()

(Not sure entirely how this is handled by graphql-spring-boot but would imagine there is some way to @Autowire or otherwise configure the GraphQL instance? InstrumentationProvider from graphql-java-servlet might be what you’d use to do this)

Okay so after spending some 20 hours on this thing . Finally figured it out .
Required a lot of trial and error . Also we most definitely dont have proper/enough documenting for this .

Although my solution answers this question , It answers an extra question as well, which is -> How do I make extensions dynamic for every query ?

Meaning , all 3 (data, errors and extensions) would be dynamic in the response .

My use case – How do I add logs for each query in extensions ?

Answer – (Note – The code below contains lombok annotations)

Step 1 – Creating your Instrumentation Class and the Instrumentation State Class .

//This is the state class . 
//State is the first thing created whenever a graphql query is run . 
//We will embed our Logs object here .
@Builder
class LogInstrumentationState implements InstrumentationState {
    @Getter
    @Setter
    public LogsDto logsDto;
}


//This is the instrumentation class that will be used to create the graphql object .
//The overridden methods are different stages in the graphql query execution 
@Builder
public class LogsInstrumentation extends SimpleInstrumentation {

  //First stage in graphql query execution .
  //Object for our custom state class object is created here . 
  @Override
  public InstrumentationState createState() {
    return LogInstrumentationState.builder().build();
  }

   //Second Stage in graphql query execution
   //Reference of initialized Logs object in the main code flow is passed here . 
   //This reference is stored in our custom state class's object .
  @Override
  public ExecutionInput instrumentExecutionInput(ExecutionInput executionInput,
                                                 InstrumentationExecutionParameters parameters) {
    LogsDto logsDto = (LogsDto) executionInput.getExtensions().get("logs");

    LogInstrumentationState logInstrumentationState = parameters.getInstrumentationState();
    logInstrumentationState.setLogsDto(logsDto);
    return super.instrumentExecutionInput(executionInput, parameters);
  }

  //This is the last stage in the graphql query execution .
  //Logs are taken from the custom container and added into extensions . 
  @Override
  public CompletableFuture<ExecutionResult> instrumentExecutionResult(
      ExecutionResult executionResult, InstrumentationExecutionParameters parameters) {

    Map<Object, Object> newExtensionMap = getExtensionsMap(executionResult,parameters);

    return CompletableFuture.completedFuture(
        new ExecutionResultImpl(
            executionResult.getData(),
            executionResult.getErrors(),
            newExtensionMap));
  }

  //Helper function
  public Map<Object, Object> getExtensionsMap(ExecutionResult executionResult, InstrumentationExecutionParameters parameters) {
    Map<Object, Object> currentExt = executionResult.getExtensions();
    Map<Object, Object> newExtensionMap = new LinkedHashMap<>();
    newExtensionMap.putAll(currentExt == null ? Collections.emptyMap() : currentExt);
    LogsDto logsDto =
        ((LogInstrumentationState)parameters.getInstrumentationState()).getLogsDto();
    newExtensionMap.put(ControllerConstants.LOGS, logsDto);
    return newExtensionMap;
  }

}

Step 2 – Creating graphql object –

GraphQL graphQl = GraphQL.newGraphQL(graphqlSchema).instrumentation(LogsInstrumentation.builder().build())
        .build();

Step 3 – Creating executing input . This is where you pass the dynamic Log object into the LogsInstrumentation class .

var executionInput = ExecutionInput.newExecutionInput()
          .query(...)
          .variables(...)
          .operationName(...)
          .extensions(Map.of("logs",logsDto))
          .dataLoaderRegistry(...)
          .graphQLContext(graphqlContext).build();


      ExecutionResult executionResult = graphQl.execute(executionInput);

Step 4 – This is how you get your extensions after your query has completed .

Map<Object, Object> extensions = executionResult.getExtensions();
LogsDto logsDto = (LogsDto) extensions.get("logs");

My source

Starting with graphql-java 21.0, you can set extensions directly on the DataFetcherResult: https://github.com/graphql-java/graphql-java/blob/v21.0/src/main/java/graphql/execution/DataFetcherResult.java#L205

It was added here. If you’re stuck on an older version of graphql-java, you can get it from the DataFetchingEnvironment like so:

ExtensionsBuilder extensionsBuilder = dfe.getGraphQlContext().get<ExtensionsBuilder>(ExtensionsBuilder.class);
if (extensionsBuilder != null) {
    extensionsBuilder.addValue("yourKey", yourValue);
}

you can implement GraphQLError where extra error properties can be added.