From 2652a57e88e4f8a6a3c1819225cb1bc882763893 Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Fri, 12 Apr 2024 15:32:07 +0100 Subject: [PATCH] Enable use of single, specific ContextSnapshotFactory instance Closes gh-919 --- .../method/InvocableHandlerMethodSupport.java | 6 +- .../ContextDataFetcherDecorator.java | 22 ++-- .../ContextSnapshotFactoryHelper.java | 117 ++++++++++++++++++ .../DataFetcherExceptionResolverAdapter.java | 5 +- .../execution/DefaultBatchLoaderRegistry.java | 7 +- .../DefaultExecutionGraphQlService.java | 6 +- .../ExceptionResolversExceptionHandler.java | 5 +- .../SubscriptionExceptionResolverAdapter.java | 5 +- .../DefaultWebGraphQlHandlerBuilder.java | 22 +++- .../graphql/server/WebGraphQlHandler.java | 20 ++- .../webmvc/GraphQlWebSocketHandler.java | 10 +- .../server/WebGraphQlHandlerTests.java | 62 ++++++++++ 12 files changed, 247 insertions(+), 40 deletions(-) create mode 100644 spring-graphql/src/main/java/org/springframework/graphql/execution/ContextSnapshotFactoryHelper.java diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java index 9a78e3cf..b49d24df 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java @@ -26,12 +26,12 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import graphql.GraphQLContext; -import io.micrometer.context.ContextSnapshotFactory; import reactor.core.publisher.Mono; import org.springframework.core.CoroutinesUtils; import org.springframework.core.KotlinDetector; import org.springframework.data.util.KotlinReflectionUtils; +import org.springframework.graphql.execution.ContextSnapshotFactoryHelper; import org.springframework.lang.Nullable; import org.springframework.util.Assert; @@ -46,8 +46,6 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod { private static final Object NO_VALUE = new Object(); - private static final ContextSnapshotFactory SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); - private final boolean hasCallableReturnValue; @@ -131,7 +129,7 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod { return CompletableFuture.supplyAsync( () -> { try { - return SNAPSHOT_FACTORY.captureFrom(graphQLContext).wrap((Callable) result).call(); + return ContextSnapshotFactoryHelper.captureFrom(graphQLContext).wrap((Callable) result).call(); } catch (Exception ex) { throw new IllegalStateException( diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java index db0e7472..1a3194d9 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java @@ -59,7 +59,6 @@ final class ContextDataFetcherDecorator implements DataFetcher { private final SubscriptionExceptionResolver subscriptionExceptionResolver; - private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); private ContextDataFetcherDecorator( DataFetcher delegate, boolean subscription, @@ -72,17 +71,18 @@ final class ContextDataFetcherDecorator implements DataFetcher { this.subscriptionExceptionResolver = subscriptionExceptionResolver; } - @Override - public Object get(DataFetchingEnvironment environment) throws Exception { - ContextSnapshot snapshot; - if (environment.getLocalContext() instanceof GraphQLContext localContext) { - snapshot = this.snapshotFactory.captureFrom(environment.getGraphQlContext(), localContext); - } - else { - snapshot = this.snapshotFactory.captureFrom(environment.getGraphQlContext()); - } - Object value = snapshot.wrap(() -> this.delegate.get(environment)).call(); + @Override + public Object get(DataFetchingEnvironment env) throws Exception { + + GraphQLContext graphQlContext = env.getGraphQlContext(); + ContextSnapshotFactory snapshotFactory = ContextSnapshotFactoryHelper.getInstance(graphQlContext); + + ContextSnapshot snapshot = (env.getLocalContext() instanceof GraphQLContext localContext) ? + snapshotFactory.captureFrom(graphQlContext, localContext) : + snapshotFactory.captureFrom(graphQlContext); + + Object value = snapshot.wrap(() -> this.delegate.get(env)).call(); if (this.subscription) { Assert.state(value instanceof Publisher, "Expected Publisher for a subscription"); diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextSnapshotFactoryHelper.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextSnapshotFactoryHelper.java new file mode 100644 index 00000000..ce3af522 --- /dev/null +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ContextSnapshotFactoryHelper.java @@ -0,0 +1,117 @@ +/* + * Copyright 2002-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.graphql.execution; + +import graphql.GraphQLContext; +import io.micrometer.context.ContextSnapshot; +import io.micrometer.context.ContextSnapshotFactory; +import reactor.util.context.Context; +import reactor.util.context.ContextView; + +import org.springframework.lang.Nullable; + +/** + * Helper to use a single {@link ContextSnapshotFactory} instance by saving and + * obtaining it to and from Reactor and GraphQL contexts. + * + * @author Rossen Stoyanchev + * @since 1.3 + */ +public abstract class ContextSnapshotFactoryHelper { + + private static final ContextSnapshotFactory sharedInstance = ContextSnapshotFactory.builder().build(); + + private static final String CONTEXT_SNAPSHOT_FACTORY_KEY = ContextSnapshotFactoryHelper.class.getName() + ".KEY"; + + + /** + * Select a {@code ContextSnapshotFactory} instance to use, either the one + * passed in if it is not {@code null}, or a shared, static instance. + * @param factory the candidate factory instance to use if not {@code null} + * @return the instance to use + */ + public static ContextSnapshotFactory selectInstance(@Nullable ContextSnapshotFactory factory) { + if (factory != null) { + return factory; + } + return sharedInstance; + } + + /** + * Save the {@code ContextSnapshotFactory} in the given {@link Context}. + * @param factory the instance to save + * @param context the context to save the instance to + * @return a new context with the saved instance + */ + public static Context saveInstance(ContextSnapshotFactory factory, Context context) { + return context.put(CONTEXT_SNAPSHOT_FACTORY_KEY, factory); + } + + /** + * Save the {@code ContextSnapshotFactory} in the given {@link Context}. + * @param factory the instance to save + * @param context the context to save the instance to + */ + public static void saveInstance(ContextSnapshotFactory factory, GraphQLContext context) { + context.put(CONTEXT_SNAPSHOT_FACTORY_KEY, factory); + } + + /** + * Access the {@code ContextSnapshotFactory} from the given {@link ContextView} + * or return a shared, static instance. + * @param contextView the context where the instance is saved + * @return the instance to use + */ + public static ContextSnapshotFactory getInstance(ContextView contextView) { + ContextSnapshotFactory factory = contextView.getOrDefault(CONTEXT_SNAPSHOT_FACTORY_KEY, null); + return selectInstance(factory); + } + + /** + * Access the {@code ContextSnapshotFactory} from the given {@link GraphQLContext} + * or return a shared, static instance. + * @param context the context where the instance is saved + * @return the instance to use + */ + public static ContextSnapshotFactory getInstance(GraphQLContext context) { + ContextSnapshotFactory factory = context.get(CONTEXT_SNAPSHOT_FACTORY_KEY); + return selectInstance(factory); + } + + /** + * Shortcut to obtain the {@code ContextSnapshotFactory} instance, and to + * capture from the given {@link ContextView}. + * @param contextView the context to capture from + * @return a snapshot from the capture + */ + public static ContextSnapshot captureFrom(ContextView contextView) { + ContextSnapshotFactory factory = getInstance(contextView); + return selectInstance(factory).captureFrom(contextView); + } + + /** + * Shortcut to obtain the {@code ContextSnapshotFactory} instance, and to + * capture from the given {@link GraphQLContext}. + * @param context the context to capture from + * @return a snapshot from the capture + */ + public static ContextSnapshot captureFrom(GraphQLContext context) { + ContextSnapshotFactory factory = getInstance(context); + return selectInstance(factory).captureFrom(context); + } + +} diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java index 07f1c27a..6e95db09 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java @@ -22,7 +22,6 @@ import java.util.function.BiFunction; import graphql.GraphQLError; import graphql.schema.DataFetchingEnvironment; -import io.micrometer.context.ContextSnapshotFactory; import io.micrometer.context.ThreadLocalAccessor; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -53,8 +52,6 @@ public abstract class DataFetcherExceptionResolverAdapter implements DataFetcher protected final Log logger = LogFactory.getLog(getClass()); - protected final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); - private boolean threadLocalContextAware; @@ -101,7 +98,7 @@ public abstract class DataFetcherExceptionResolverAdapter implements DataFetcher return resolveToMultipleErrors(exception, env); } try { - return this.snapshotFactory.captureFrom(env.getGraphQlContext()) + return ContextSnapshotFactoryHelper.captureFrom(env.getGraphQlContext()) .wrap(() -> resolveToMultipleErrors(exception, env)) .call(); } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java index 1c204dfe..f938a047 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java @@ -28,7 +28,6 @@ import java.util.function.Supplier; import graphql.GraphQLContext; import io.micrometer.context.ContextSnapshot; -import io.micrometer.context.ContextSnapshotFactory; import org.dataloader.BatchLoaderContextProvider; import org.dataloader.BatchLoaderEnvironment; import org.dataloader.BatchLoaderWithContext; @@ -54,8 +53,6 @@ import org.springframework.util.StringUtils; */ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry { - private static final ContextSnapshotFactory SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); - private final List> loaders = new ArrayList<>(); private final List> mappedLoaders = new ArrayList<>(); @@ -231,7 +228,7 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry { @Override public CompletionStage> load(List keys, BatchLoaderEnvironment environment) { GraphQLContext graphQLContext = environment.getContext(); - ContextSnapshot snapshot = SNAPSHOT_FACTORY.captureFrom(graphQLContext); + ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(graphQLContext); try { return snapshot.wrap(() -> this.loader.apply(keys, environment) @@ -279,7 +276,7 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry { @Override public CompletionStage> load(Set keys, BatchLoaderEnvironment environment) { GraphQLContext graphQLContext = environment.getContext(); - ContextSnapshot snapshot = SNAPSHOT_FACTORY.captureFrom(graphQLContext); + ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(graphQLContext); try { return snapshot.wrap(() -> this.loader.apply(keys, environment) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java index b88ba2e0..68da48b8 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java @@ -46,7 +46,6 @@ public class DefaultExecutionGraphQlService implements ExecutionGraphQlService { private static final BiFunction RESET_EXECUTION_ID_CONFIGURER = (executionInput, builder) -> builder.executionId(null).build(); - private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); private final GraphQlSource graphQlSource; @@ -90,7 +89,10 @@ public class DefaultExecutionGraphQlService implements ExecutionGraphQlService { ExecutionInput executionInput = request.toExecutionInput(); - this.snapshotFactory.captureFrom(contextView).updateContext(executionInput.getGraphQLContext()); + ContextSnapshotFactory factory = ContextSnapshotFactoryHelper.getInstance(contextView); + GraphQLContext graphQLContext = executionInput.getGraphQLContext(); + ContextSnapshotFactoryHelper.saveInstance(factory, graphQLContext); + factory.captureFrom(contextView).updateContext(graphQLContext); ExecutionInput updatedExecutionInput = (this.hasDataLoaderRegistrations ? registerDataLoaders(executionInput) : executionInput); diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java index e5dd044d..d62695a7 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java @@ -29,7 +29,6 @@ import graphql.execution.DataFetcherExceptionHandlerResult; import graphql.execution.ExecutionId; import graphql.schema.DataFetchingEnvironment; import io.micrometer.context.ContextSnapshot; -import io.micrometer.context.ContextSnapshotFactory; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import reactor.core.publisher.Flux; @@ -47,8 +46,6 @@ class ExceptionResolversExceptionHandler implements DataFetcherExceptionHandler private static final Log logger = LogFactory.getLog(ExceptionResolversExceptionHandler.class); - private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); - private final List resolvers; /** @@ -65,7 +62,7 @@ class ExceptionResolversExceptionHandler implements DataFetcherExceptionHandler public CompletableFuture handleException(DataFetcherExceptionHandlerParameters params) { Throwable exception = unwrapException(params); DataFetchingEnvironment env = params.getDataFetchingEnvironment(); - ContextSnapshot snapshot = this.snapshotFactory.captureFrom(env.getGraphQlContext()); + ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(env.getGraphQlContext()); try { return Flux.fromIterable(this.resolvers) .flatMap((resolver) -> resolver.resolveException(exception, env)) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java index 37f2078d..ff87887a 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/SubscriptionExceptionResolverAdapter.java @@ -22,7 +22,6 @@ import java.util.function.Function; import graphql.GraphQLError; import io.micrometer.context.ContextSnapshot; -import io.micrometer.context.ContextSnapshotFactory; import io.micrometer.context.ThreadLocalAccessor; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -51,8 +50,6 @@ public abstract class SubscriptionExceptionResolverAdapter implements Subscripti protected final Log logger = LogFactory.getLog(getClass()); - protected final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); - private boolean threadLocalContextAware; @@ -86,7 +83,7 @@ public abstract class SubscriptionExceptionResolverAdapter implements Subscripti public final Mono> resolveException(Throwable exception) { if (this.threadLocalContextAware) { return Mono.deferContextual((contextView) -> { - ContextSnapshot snapshot = this.snapshotFactory.captureFrom(contextView); + ContextSnapshot snapshot = ContextSnapshotFactoryHelper.captureFrom(contextView); try { List errors = snapshot.wrap(() -> resolveToMultipleErrors(exception)).call(); return Mono.justOrEmpty(errors); diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java b/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java index e00f57d6..f6bb522b 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/server/DefaultWebGraphQlHandlerBuilder.java @@ -25,6 +25,7 @@ import io.micrometer.context.ContextSnapshotFactory; import reactor.core.publisher.Mono; import org.springframework.graphql.ExecutionGraphQlService; +import org.springframework.graphql.execution.ContextSnapshotFactoryHelper; import org.springframework.graphql.server.WebGraphQlInterceptor.Chain; import org.springframework.lang.Nullable; import org.springframework.util.Assert; @@ -41,6 +42,9 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { private final List interceptors = new ArrayList<>(); + @Nullable + private ContextSnapshotFactory snapshotFactory; + @Nullable private WebSocketGraphQlInterceptor webSocketInterceptor; @@ -68,10 +72,16 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { return this; } + @Override + public WebGraphQlHandler.Builder contextSnapshotFactory(ContextSnapshotFactory snapshotFactory) { + this.snapshotFactory = snapshotFactory; + return this; + } + @Override public WebGraphQlHandler build() { - ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); + ContextSnapshotFactory snapshotFactory = ContextSnapshotFactoryHelper.selectInstance(this.snapshotFactory); Chain endOfChain = (request) -> this.service.execute(request).map(WebGraphQlResponse::new); @@ -88,10 +98,18 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { DefaultWebGraphQlHandlerBuilder.this.webSocketInterceptor : new WebSocketGraphQlInterceptor() { }; } + @Override + public ContextSnapshotFactory contextSnapshotFactory() { + return snapshotFactory; + } + @Override public Mono handleRequest(WebGraphQlRequest request) { ContextSnapshot snapshot = snapshotFactory.captureAll(); - return executionChain.next(request).contextWrite(snapshot::updateContext); + return executionChain.next(request).contextWrite((context) -> { + context = ContextSnapshotFactoryHelper.saveInstance(snapshotFactory, context); + return snapshot.updateContext(context); + }); } }; } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlHandler.java index fbdfaf85..45b68eb0 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/server/WebGraphQlHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2022 the original author or authors. + * Copyright 2002-2024 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,6 +18,7 @@ package org.springframework.graphql.server; import java.util.List; +import io.micrometer.context.ContextSnapshotFactory; import reactor.core.publisher.Mono; import org.springframework.graphql.ExecutionGraphQlService; @@ -39,6 +40,13 @@ public interface WebGraphQlHandler { */ WebSocketGraphQlInterceptor getWebSocketInterceptor(); + /** + * Return the {@link WebGraphQlHandler.Builder#contextSnapshotFactory configured} + * {@code ContextSnapshotFactory} instance to use. + * @since 1.3 + */ + ContextSnapshotFactory contextSnapshotFactory(); + /** * Execute the given request and return the response. * @param request the request to execute @@ -86,6 +94,16 @@ public interface WebGraphQlHandler { */ Builder interceptors(List interceptors); + /** + * Configure the {@link ContextSnapshotFactory} instance to use for + * context propagation of {@code ThreadLocal}, and Reactor context + * values from the transport layer to the GraphQL execution layer. + * If not set, then a default instance is used. + * @param snapshotFactory the factory to use + * @since 1.3 + */ + Builder contextSnapshotFactory(ContextSnapshotFactory snapshotFactory); + /** * Build the {@link WebGraphQlHandler} instance. * @return the built WebGraphQlHandler diff --git a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java index 1f486827..a03673b2 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/server/webmvc/GraphQlWebSocketHandler.java @@ -141,7 +141,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub Assert.notNull(converter, "HttpMessageConverter for JSON is required"); this.graphQlHandler = graphQlHandler; - this.contextHandshakeInterceptor = new ContextHandshakeInterceptor(); + this.contextHandshakeInterceptor = new ContextHandshakeInterceptor(graphQlHandler.contextSnapshotFactory()); this.webSocketGraphQlInterceptor = this.graphQlHandler.getWebSocketInterceptor(); this.initTimeoutDuration = connectionInitTimeout; this.converter = converter; @@ -392,14 +392,18 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub private static final String KEY = ContextSnapshot.class.getName(); - private static final ContextSnapshotFactory SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); + private final ContextSnapshotFactory snapshotFactory; + + ContextHandshakeInterceptor(ContextSnapshotFactory factory) { + this.snapshotFactory = factory; + } @Override public boolean beforeHandshake( ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map attributes) { - attributes.put(KEY, SNAPSHOT_FACTORY.captureAll()); + attributes.put(KEY, this.snapshotFactory.captureAll()); return true; } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java index ed57dd1f..acf37e86 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/server/WebGraphQlHandlerTests.java @@ -23,9 +23,12 @@ import java.util.Collections; import graphql.GraphqlErrorBuilder; import graphql.schema.DataFetcher; import io.micrometer.context.ContextRegistry; +import io.micrometer.context.ContextSnapshot; +import io.micrometer.context.ContextSnapshotFactory; import org.junit.jupiter.api.Test; import reactor.core.publisher.Mono; +import org.springframework.graphql.ExecutionGraphQlService; import org.springframework.graphql.GraphQlSetup; import org.springframework.graphql.ResponseHelper; import org.springframework.graphql.TestThreadLocalAccessor; @@ -141,4 +144,63 @@ public class WebGraphQlHandlerTests { } } + @Test + void contextSnapshotFactoryInstance() { + + DataFetcherExceptionResolver exceptionResolver = + (ex, env) -> Mono.deferContextual((view) -> Mono.just(Collections.singletonList( + GraphqlErrorBuilder.newError(env) + .message("Resolved error: " + ex.getMessage() + ", name=" + view.get("name")) + .errorType(ErrorType.BAD_REQUEST) + .build()))); + + ExecutionGraphQlService service = this.graphQlSetup + .queryFetcher("greeting", this.errorDataFetcher) + .exceptionResolver(exceptionResolver) + .toGraphQlService(); + + InvocationCountingContextSnapshotFactory factory = new InvocationCountingContextSnapshotFactory(); + WebGraphQlHandler handler = WebGraphQlHandler.builder(service).contextSnapshotFactory(factory).build(); + + Mono responseMono = + handler.handleRequest(webInput).contextWrite((context) -> context.put("name", "007")); + + ResponseHelper response = ResponseHelper.forResponse(responseMono); + assertThat(response.errorCount()).isEqualTo(1); + assertThat(response.error(0).message()).isEqualTo("Resolved error: Invalid greeting, name=007"); + + assertThat(factory.captureAllCount).isEqualTo(1); + assertThat(factory.captureFromCount) + .as("One or more of the following did not use the configured instance: " + + "DefaultExecutionService, ContextDataFetcher, or ExceptionResolversExceptionHandler.") + .isEqualTo(3); + } + + + private static class InvocationCountingContextSnapshotFactory implements ContextSnapshotFactory { + + private final ContextSnapshotFactory delegate = ContextSnapshotFactory.builder().build(); + + private int captureAllCount; + + private int captureFromCount; + + @Override + public ContextSnapshot captureAll(Object... contexts) { + this.captureAllCount++; + return this.delegate.captureAll(contexts); + } + + @Override + public ContextSnapshot captureFrom(Object... contexts) { + this.captureFromCount++; + return this.delegate.captureFrom(contexts); + } + + @Override + public ContextSnapshot.Scope setThreadLocalsFrom(Object context, String... keys) { + return this.delegate.setThreadLocalsFrom(context, keys); + } + } + }