Skip to content

Commit 58bd58b

Browse files
koenpuntbclozel
authored andcommitted
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
1 parent 29fb329 commit 58bd58b

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

spring-graphql/src/main/java/org/springframework/graphql/observation/GraphQlObservationInstrumentation.java

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -155,13 +155,13 @@ public DataFetcher<?> instrumentDataFetcher(DataFetcher<?> dataFetcher,
155155
throw new CompletionException(error);
156156
}
157157
dataFetcherObservation.stop();
158-
return wrapAsDataFetcherResult(result, dataFetcherObservation);
158+
return wrapAsDataFetcherResult(result, dataFetcherObservation, environment.getLocalContext());
159159
});
160160
}
161161
else {
162162
observationContext.setValue(value);
163163
dataFetcherObservation.stop();
164-
return wrapAsDataFetcherResult(value, dataFetcherObservation);
164+
return wrapAsDataFetcherResult(value, dataFetcherObservation, environment.getLocalContext());
165165
}
166166
}
167167
catch (Throwable throwable) {
@@ -187,7 +187,8 @@ private static Observation getCurrentObservation(DataFetchingEnvironment environ
187187
return currentObservation;
188188
}
189189

190-
private static DataFetcherResult<?> wrapAsDataFetcherResult(Object value, Observation dataFetcherObservation) {
190+
private static DataFetcherResult<?> wrapAsDataFetcherResult(Object value, Observation dataFetcherObservation,
191+
@Nullable GraphQLContext dataFetcherLocalContext) {
191192
if (value instanceof DataFetcherResult<?> result) {
192193
if (result.getLocalContext() == null) {
193194
return result.transform(builder -> builder.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build()));
@@ -201,9 +202,11 @@ else if (result.getLocalContext() instanceof GraphQLContext) {
201202
return result;
202203
}
203204
else {
205+
GraphQLContext localContext = dataFetcherLocalContext == null ?
206+
GraphQLContext.getDefault() : dataFetcherLocalContext;
204207
return DataFetcherResult.newResult()
205208
.data(value)
206-
.localContext(GraphQLContext.newContext().of(ObservationThreadLocalAccessor.KEY, dataFetcherObservation).build())
209+
.localContext(localContext.put(ObservationThreadLocalAccessor.KEY, dataFetcherObservation))
207210
.build();
208211
}
209212

0 commit comments

Comments
 (0)