From 58bd58bef06b113ac5e2da4e51f4201c2edaf12a Mon Sep 17 00:00:00 2001 From: Koen Punt Date: Thu, 20 Jul 2023 18:46:50 +0200 Subject: [PATCH] Inherit local context from environment in instrumentation Prior to this commit, a `DataFetcher` instrumented by the `GraphQlObservationInstrumentation` would incorrectly overwrite the local context when: * the `DataFetcher` returns a value object (i.e. not a `DataFetcherResult`) * the given `DatFetchingEnvironment` has an existing local context with values For this case, the instrumentation would create a new local context but would not inherit from the existing local context. See gh-761 --- .../GraphQlObservationInstrumentation.java | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java b/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java index ac33db1d..47567553 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java @@ -155,13 +155,13 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen throw new CompletionException(error); } dataFetcherObservation.stop(); - return wrapAsDataFetcherResult(result, dataFetcherObservation); + return wrapAsDataFetcherResult(result, dataFetcherObservation, environment.getLocalContext()); }); } else { observationContext.setValue(value); dataFetcherObservation.stop(); - return wrapAsDataFetcherResult(value, dataFetcherObservation); + return wrapAsDataFetcherResult(value, dataFetcherObservation, environment.getLocalContext()); } } catch (Throwable throwable) { @@ -187,7 +187,8 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen return currentObservation; } - private static DataFetcherResult wrapAsDataFetcherResult(Object value, Observation dataFetcherObservation) { + private static DataFetcherResult wrapAsDataFetcherResult(Object value, Observation dataFetcherObservation, + @Nullable GraphQLContext dataFetcherLocalContext) { if (value instanceof DataFetcherResult result) { if (result.getLocalContext() == null) { return result.transform(builder -> builder.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build())); @@ -201,9 +202,11 @@ public class GraphQlObservationInstrumentation extends SimplePerformantInstrumen return result; } else { + GraphQLContext localContext = dataFetcherLocalContext == null ? + GraphQLContext.getDefault() : dataFetcherLocalContext; return DataFetcherResult.newResult() .data(value) - .localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build()) + .localContext(localContext.put(ObservationThreadLocalAccessor.KEY, dataFetcherObservation)) .build(); }