From 01ace235188e674cc692c913c75271c32d7a24aa Mon Sep 17 00:00:00 2001 From: Antoine Boyer Date: Thu, 8 Jan 2026 16:34:57 -0800 Subject: [PATCH] Add FieldFetchingInstrumentationContext.onExceptionHandled for hooking after exception handling --- .../graphql/execution/ExecutionStrategy.java | 21 ++- .../FieldFetchingInstrumentationContext.java | 10 ++ .../InstrumentationTest.groovy | 166 ++++++++++++++++++ 3 files changed, 189 insertions(+), 8 deletions(-) diff --git a/src/main/java/graphql/execution/ExecutionStrategy.java b/src/main/java/graphql/execution/ExecutionStrategy.java index 06d3b644b1..78ad4280b8 100644 --- a/src/main/java/graphql/execution/ExecutionStrategy.java +++ b/src/main/java/graphql/execution/ExecutionStrategy.java @@ -489,12 +489,17 @@ private Object fetchField(GraphQLFieldDefinition fieldDef, ExecutionContext exec CompletableFuture> handleCF = engineRunningState.handle(fetchedValue, (result, exception) -> { // because we added an artificial CF, we need to unwrap the exception - fetchCtx.onCompleted(result, exception); - exception = engineRunningState.possibleCancellation(exception); - - if (exception != null) { - return handleFetchingException(dataFetchingEnvironment.get(), parameters, exception); + Throwable possibleWrappedException = engineRunningState.possibleCancellation(exception); + + if (possibleWrappedException != null) { + CompletableFuture> handledExceptionResult = handleFetchingException(dataFetchingEnvironment.get(), parameters, possibleWrappedException); + return handledExceptionResult.thenApply( handledResult -> { + fetchCtx.onExceptionHandled(handledResult); + fetchCtx.onCompleted(result, exception); + return handledResult; + }); } else { + fetchCtx.onCompleted(result, exception); // we can simply return the fetched value CF and avoid a allocation return fetchedValue; } @@ -578,7 +583,7 @@ private void addExtensionsIfPresent(ExecutionContext executionContext, DataFetch } } - protected CompletableFuture handleFetchingException( + protected CompletableFuture> handleFetchingException( DataFetchingEnvironment environment, ExecutionStrategyParameters parameters, Throwable e @@ -599,10 +604,10 @@ protected CompletableFuture handleFetchingException( } } - private CompletableFuture asyncHandleException(DataFetcherExceptionHandler handler, DataFetcherExceptionHandlerParameters handlerParameters) { + private CompletableFuture> asyncHandleException(DataFetcherExceptionHandler handler, DataFetcherExceptionHandlerParameters handlerParameters) { //noinspection unchecked return handler.handleException(handlerParameters).thenApply( - handlerResult -> (T) DataFetcherResult.newResult().errors(handlerResult.getErrors()).build() + handlerResult -> (DataFetcherResult) DataFetcherResult.newResult().errors(handlerResult.getErrors()).build() ); } diff --git a/src/main/java/graphql/execution/instrumentation/FieldFetchingInstrumentationContext.java b/src/main/java/graphql/execution/instrumentation/FieldFetchingInstrumentationContext.java index 38984c6f92..f6ff09beca 100644 --- a/src/main/java/graphql/execution/instrumentation/FieldFetchingInstrumentationContext.java +++ b/src/main/java/graphql/execution/instrumentation/FieldFetchingInstrumentationContext.java @@ -2,6 +2,7 @@ import graphql.Internal; import graphql.PublicSpi; +import graphql.execution.DataFetcherResult; import graphql.execution.instrumentation.parameters.InstrumentationFieldFetchParameters; import org.jspecify.annotations.NonNull; import org.jspecify.annotations.Nullable; @@ -26,6 +27,15 @@ public interface FieldFetchingInstrumentationContext extends InstrumentationCont default void onFetchedValue(Object fetchedValue) { } + /** + * This is called back after any {@link graphql.execution.DataFetcherExceptionHandler}) has run on any exception raised + * during a {@link graphql.schema.DataFetcher} invocation. This allows to see the final {@link DataFetcherResult} + * that will be used when performing the complete step. + * @param dataFetcherResult the final {@link DataFetcherResult} after the exception handler has run + */ + default void onExceptionHandled(DataFetcherResult dataFetcherResult) { + } + @Internal FieldFetchingInstrumentationContext NOOP = new FieldFetchingInstrumentationContext() { @Override diff --git a/src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy b/src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy index 43767d9348..ca52f8cf84 100644 --- a/src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy +++ b/src/test/groovy/graphql/execution/instrumentation/InstrumentationTest.groovy @@ -1,11 +1,16 @@ package graphql.execution.instrumentation +import graphql.ErrorType import graphql.ExecutionInput import graphql.ExecutionResult import graphql.GraphQL +import graphql.GraphqlErrorBuilder +import graphql.GraphqlErrorBuilderTest import graphql.StarWarsSchema import graphql.TestUtil import graphql.execution.AsyncExecutionStrategy +import graphql.execution.DataFetcherExceptionHandlerResult +import graphql.execution.DataFetcherResult import graphql.execution.instrumentation.parameters.InstrumentationCreateStateParameters import graphql.execution.instrumentation.parameters.InstrumentationExecutionParameters import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters @@ -23,6 +28,8 @@ import spock.lang.Specification import java.util.concurrent.CompletableFuture import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicBoolean +import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicLong class InstrumentationTest extends Specification { @@ -152,6 +159,165 @@ class InstrumentationTest extends Specification { instrumentation.throwableList[0].getMessage() == "DF BANG!" } + def "field fetch will instrument exceptions correctly - includes exception handling with onExceptionHandled"() { + + given: + + def query = """ + { + hero { + id + } + } + """ + + def instrumentation = new LegacyTestingInstrumentation() { + def onHandledCalled = false + def onCompletedCalled = false + def onDispatchedCalled = false + + @Override + DataFetcher instrumentDataFetcher(DataFetcher dataFetcher, InstrumentationFieldFetchParameters parameters, InstrumentationState state) { + return new DataFetcher() { + @Override + Object get(DataFetchingEnvironment environment) { + throw new RuntimeException("DF BANG!") + } + } + } + + @Override + FieldFetchingInstrumentationContext beginFieldFetching(InstrumentationFieldFetchParameters parameters, InstrumentationState state) { + return new FieldFetchingInstrumentationContext() { + @Override + void onDispatched() { + onDispatchedCalled = true + } + + @Override + void onCompleted(Object result, Throwable t) { + onCompletedCalled = true + } + + @Override + void onExceptionHandled(DataFetcherResult dataFetcherResult) { + onHandledCalled = true + } + } + } + } + + def graphQL = GraphQL + .newGraphQL(StarWarsSchema.starWarsSchema) + .defaultDataFetcherExceptionHandler { it -> + // catch all exceptions and transform to graphql error with a prefixed message + return CompletableFuture.completedFuture( + DataFetcherExceptionHandlerResult.newResult(GraphqlErrorBuilder.newError() + .errorType(ErrorType.DataFetchingException) + .message("Handled " + it.exception.message) + .path(it.path) + .build()) + .build()) + } + .instrumentation(instrumentation) + .build() + + when: + def resp = graphQL.execute(query) + + then: "exception handler turned the exception into a graphql error and message prefixed with Handled" + resp.errors.size() == 1 + resp.errors[0].message == "Handled DF BANG!" + + and: "all instrumentation methods were called" + instrumentation.onDispatchedCalled == true + instrumentation.onCompletedCalled == true + instrumentation.onHandledCalled == true + } + + + def "field fetch verify order and call of all methods"() { + + given: + + def query = """ + { + hero { + id + } + } + """ + + def metric = [] + def instrumentation = new SimplePerformantInstrumentation() { + def timeElapsed = new AtomicInteger() + + @Override + DataFetcher instrumentDataFetcher(DataFetcher dataFetcher, InstrumentationFieldFetchParameters parameters, InstrumentationState state) { + return new DataFetcher() { + @Override + Object get(DataFetchingEnvironment environment) { + // simulate latency + timeElapsed.addAndGet(50) + throw new RuntimeException("DF BANG!") + } + } + } + + @Override + FieldFetchingInstrumentationContext beginFieldFetching(InstrumentationFieldFetchParameters parameters, InstrumentationState state) { + return new FieldFetchingInstrumentationContext() { + def start = 0 + def duration = 0 + def hasError = false + + + @Override + void onDispatched() { + start = 1 + } + + @Override + void onCompleted(Object result, Throwable t) { + duration = timeElapsed.get() - start + metric = [duration, hasError] + } + + @Override + void onExceptionHandled(DataFetcherResult dataFetcherResult) { + hasError = dataFetcherResult.errors != null && !dataFetcherResult.errors.isEmpty() + && dataFetcherResult.errors.any { it.message.contains("Handled") } + } + } + } + } + + def graphQL = GraphQL + .newGraphQL(StarWarsSchema.starWarsSchema) + .defaultDataFetcherExceptionHandler { it -> + // catch all exceptions and transform to graphql error with a prefixed message + return CompletableFuture.completedFuture( + DataFetcherExceptionHandlerResult.newResult(GraphqlErrorBuilder.newError() + .errorType(ErrorType.DataFetchingException) + .message("Handled " + it.exception.message) + .path(it.path) + .build()) + .build()) + } + .instrumentation(instrumentation) + .build() + + when: + def resp = graphQL.execute(query) + + then: "exception handler turned the exception into a graphql error and prefixed its message with 'Handled'" + resp.errors.size() == 1 + resp.errors[0].message == "Handled DF BANG!" + + and: "metric was captured i.e all instrumentation methods were called in the right order" + metric == [49, true] + } + /** * This uses a stop and go pattern and multiple threads. Each time * the execution strategy is invoked, the data fetchers are held