From dcaa1511a11c4fc00e7b6dd774e12126617364f3 Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Wed, 13 Mar 2024 13:30:33 +0000 Subject: [PATCH] Configurable ContextSnapshotFactory instance Closes gh-919 --- .../data/federation/EntityHandlerMethod.java | 14 ++++++- .../federation/FederationSchemaFactory.java | 4 +- .../method/InvocableHandlerMethodSupport.java | 30 ++++++++++--- .../AnnotatedControllerConfigurer.java | 14 +++++-- .../AnnotatedControllerDetectionSupport.java | 19 +++++++++ .../AnnotatedControllerExceptionResolver.java | 6 +-- .../support/BatchLoaderHandlerMethod.java | 25 ++++++++++- .../support/DataFetcherHandlerMethod.java | 30 ++++++++++--- .../DataFetcherHandlerMethodSupport.java | 24 ++++++++++- .../AbstractGraphQlSourceBuilder.java | 27 ++++++++++-- .../ContextDataFetcherDecorator.java | 30 ++++++++++--- .../DataFetcherExceptionResolver.java | 22 +++++++++- .../DataFetcherExceptionResolverAdapter.java | 19 ++++++++- .../execution/DefaultBatchLoaderRegistry.java | 42 +++++++++++++++---- .../DefaultExecutionGraphQlService.java | 22 +++++++++- .../ExceptionResolversExceptionHandler.java | 10 ++++- .../graphql/execution/GraphQlSource.java | 20 ++++++++- .../SubscriptionExceptionResolverAdapter.java | 14 ++++++- .../DefaultWebGraphQlHandlerBuilder.java | 19 ++++++++- .../graphql/server/WebGraphQlHandler.java | 29 ++++++++++++- .../webmvc/GraphQlWebSocketHandler.java | 10 +++-- ...ntextValueMethodArgumentResolverTests.java | 6 +-- .../DataFetcherHandlerMethodTests.java | 8 ++-- 23 files changed, 377 insertions(+), 67 deletions(-) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java index 1f457273..7e1b6e78 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/EntityHandlerMethod.java @@ -17,10 +17,12 @@ package org.springframework.graphql.data.federation; import java.util.Map; +import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import graphql.schema.DataFetchingEnvironment; +import io.micrometer.context.ContextSnapshotFactory; import reactor.core.publisher.Mono; import org.springframework.graphql.data.method.HandlerMethod; @@ -36,11 +38,19 @@ import org.springframework.lang.Nullable; */ final class EntityHandlerMethod extends DataFetcherHandlerMethodSupport { + /** + * Create an instance. + * @param handlerMethod the handler method + * @param resolvers the argument resolvers + * @param executor {@code Executor} to use for {@link Callable} methods + * @param snapshotFactory for context propagation with {@link Callable} methods + * @since 1.3 + */ public EntityHandlerMethod( HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers, - @Nullable Executor executor) { + @Nullable Executor executor, @Nullable ContextSnapshotFactory snapshotFactory) { - super(handlerMethod, resolvers, executor); + super(handlerMethod, resolvers, executor, snapshotFactory); } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java index 0ecaac79..1c0b865c 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/federation/FederationSchemaFactory.java @@ -87,8 +87,8 @@ public final class FederationSchemaFactory super.afterPropertiesSet(); detectHandlerMethods().forEach(info -> - this.handlerMethods.put(info.typeName(), - new EntityHandlerMethod(info.handlerMethod(), getArgumentResolvers(), getExecutor()))); + this.handlerMethods.put(info.typeName(), new EntityHandlerMethod( + info.handlerMethod(), getArgumentResolvers(), getExecutor(), getContextSnapshotFactory()))); if (this.typeResolver == null) { this.typeResolver = new ClassNameTypeResolver(); 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 9f6183a2..1d2b288a 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 @@ -46,7 +46,7 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod { private static final Object NO_VALUE = new Object(); - private static final ContextSnapshotFactory SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); + private static final ContextSnapshotFactory DEFAULT_SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); private final boolean hasCallableReturnValue; @@ -54,16 +54,35 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod { @Nullable private final Executor executor; + private final ContextSnapshotFactory snapshotFactory; + + + /** + * Create an instance for a controller method and an optional {@link Executor} + * to use for {@link Callable} return values. + * @deprecated in favor of + * {@link #InvocableHandlerMethodSupport(HandlerMethod, Executor, ContextSnapshotFactory)}. + */ + @Deprecated(since = "1.3", forRemoval = true) + protected InvocableHandlerMethodSupport(HandlerMethod handlerMethod, @Nullable Executor executor) { + this(handlerMethod, executor, null); + } /** * Create an instance. - * @param handlerMethod the controller method - * @param executor an {@link Executor} to use for {@link Callable} return values + * @param handlerMethod the handler method + * @param executor {@code Executor} to use for {@link Callable} methods + * @param snapshotFactory for context propagation with {@link Callable} methods + * @since 1.3 */ - protected InvocableHandlerMethodSupport(HandlerMethod handlerMethod, @Nullable Executor executor) { + protected InvocableHandlerMethodSupport( + HandlerMethod handlerMethod, @Nullable Executor executor, @Nullable ContextSnapshotFactory snapshotFactory) { + super(handlerMethod.createWithResolvedBean()); this.hasCallableReturnValue = getReturnType().getParameterType().equals(Callable.class); this.executor = executor; + this.snapshotFactory = (snapshotFactory != null ? snapshotFactory : DEFAULT_SNAPSHOT_FACTORY); + Assert.isTrue(!this.hasCallableReturnValue || this.executor != null, "Controller method declared with Callable return value, but no Executor configured: " + handlerMethod.getBridgedMethod().toGenericString()); @@ -130,7 +149,8 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod { return CompletableFuture.supplyAsync( () -> { try { - return SNAPSHOT_FACTORY.captureFrom(graphQLContext).wrap((Callable) result).call(); + ContextSnapshot snapshot = this.snapshotFactory.captureFrom(graphQLContext); + return snapshot.wrap((Callable) result).call(); } catch (Exception ex) { throw new IllegalStateException( diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java index d9bec726..94e5012d 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java @@ -36,6 +36,7 @@ import graphql.schema.DataFetchingEnvironment; import graphql.schema.FieldCoordinates; import graphql.schema.GraphQLCodeRegistry; import graphql.schema.idl.RuntimeWiring; +import io.micrometer.context.ContextSnapshotFactory; import org.dataloader.DataLoader; import org.reactivestreams.Publisher; import reactor.core.publisher.Flux; @@ -226,7 +227,8 @@ public class AnnotatedControllerConfigurer DataFetcher dataFetcher; if (!info.isBatchMapping()) { dataFetcher = new SchemaMappingDataFetcher( - info, getArgumentResolvers(), this.validationHelper, getExceptionResolver(), getExecutor()); + info, getArgumentResolvers(), this.validationHelper, getExceptionResolver(), + getExecutor(), getContextSnapshotFactory()); } else { dataFetcher = registerBatchLoader(info); @@ -320,7 +322,8 @@ public class AnnotatedControllerConfigurer } HandlerMethod handlerMethod = info.getHandlerMethod(); - BatchLoaderHandlerMethod invocable = new BatchLoaderHandlerMethod(handlerMethod, getExecutor()); + BatchLoaderHandlerMethod invocable = + new BatchLoaderHandlerMethod(handlerMethod, getExecutor(), getContextSnapshotFactory()); MethodParameter returnType = handlerMethod.getReturnType(); Class clazz = returnType.getParameterType(); @@ -389,15 +392,18 @@ public class AnnotatedControllerConfigurer @Nullable private final Executor executor; + @Nullable final ContextSnapshotFactory snapshotFactory; + private final boolean subscription; SchemaMappingDataFetcher( DataFetcherMappingInfo info, HandlerMethodArgumentResolverComposite argumentResolvers, @Nullable ValidationHelper helper, HandlerDataFetcherExceptionResolver exceptionResolver, - @Nullable Executor executor) { + @Nullable Executor executor, @Nullable ContextSnapshotFactory snapshotFactory) { this.mappingInfo = info; this.argumentResolvers = argumentResolvers; + this.snapshotFactory = snapshotFactory; this.methodValidationHelper = (helper != null ? helper.getValidationHelperFor(info.getHandlerMethod()) : null); @@ -431,7 +437,7 @@ public class AnnotatedControllerConfigurer DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod( getHandlerMethod(), this.argumentResolvers, this.methodValidationHelper, - this.executor, this.subscription); + this.subscription, this.executor, this.snapshotFactory); try { Object result = handlerMethod.invoke(environment); diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerDetectionSupport.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerDetectionSupport.java index 9fada65c..b8c9a8df 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerDetectionSupport.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerDetectionSupport.java @@ -29,6 +29,7 @@ import java.util.concurrent.Executor; import java.util.stream.Collectors; import graphql.schema.DataFetcher; +import io.micrometer.context.ContextSnapshotFactory; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; @@ -90,6 +91,9 @@ public abstract class AnnotatedControllerDetectionSupport implements Applicat @Nullable private Executor executor; + @Nullable + private ContextSnapshotFactory snapshotFactory; + @Nullable private HandlerMethodArgumentResolverComposite argumentResolvers; @@ -161,6 +165,21 @@ public abstract class AnnotatedControllerDetectionSupport implements Applicat return this.executor; } + /** + * Configure the {@link ContextSnapshotFactory} instance to use to establish + * {@code ThreadLocal} context for asynchronous handling with a {@link Callable} + * return value. If not set, then an instance with default settings is used. + * @since 1.3 + */ + public void setContextSnapshotFactory(@Nullable ContextSnapshotFactory snapshotFactory) { + this.snapshotFactory = snapshotFactory; + } + + @Nullable + public ContextSnapshotFactory getContextSnapshotFactory() { + return this.snapshotFactory; + } + /** * Return the configured argument resolvers. */ diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerExceptionResolver.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerExceptionResolver.java index 006ac950..69541080 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerExceptionResolver.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerExceptionResolver.java @@ -215,9 +215,9 @@ final class AnnotatedControllerExceptionResolver implements HandlerDataFetcherEx private Mono> invokeExceptionHandler( Throwable exception, DataFetchingEnvironment env, Object controllerOrAdvice, MethodHolder methodHolder) { - DataFetcherHandlerMethod exceptionHandler = new DataFetcherHandlerMethod( - new HandlerMethod(controllerOrAdvice, methodHolder.getMethod()), this.argumentResolvers, - null, null, false); + HandlerMethod handlerMethod = new HandlerMethod(controllerOrAdvice, methodHolder.getMethod()); + DataFetcherHandlerMethod exceptionHandler = + new DataFetcherHandlerMethod(handlerMethod, this.argumentResolvers, null, false, null, null); List exceptions = new ArrayList<>(); try { diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java index 0a255542..41c747ee 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2023 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. @@ -19,11 +19,13 @@ import java.security.Principal; import java.util.Arrays; import java.util.Collection; import java.util.Map; +import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import java.util.function.Function; import graphql.GraphQLContext; +import io.micrometer.context.ContextSnapshotFactory; import org.dataloader.BatchLoaderEnvironment; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -58,8 +60,27 @@ public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport { private final ParameterNameDiscoverer parameterNameDiscoverer = new DefaultParameterNameDiscoverer(); + /** + * @deprecated in favor of + * {@link #BatchLoaderHandlerMethod(HandlerMethod, Executor, ContextSnapshotFactory)}. + */ + @Deprecated(since = "1.3", forRemoval = true) public BatchLoaderHandlerMethod(HandlerMethod handlerMethod, @Nullable Executor executor) { - super(handlerMethod, executor); + this(handlerMethod, executor, null); + } + + /** + * Create an instance. + * @param handlerMethod the handler method + * @param executor {@code Executor} to use for {@link Callable} methods + * @param snapshotFactory for context propagation with {@link Callable} methods + * @since 1.3 + */ + public BatchLoaderHandlerMethod( + HandlerMethod handlerMethod, @Nullable Executor executor, + @Nullable ContextSnapshotFactory snapshotFactory) { + + super(handlerMethod, executor, snapshotFactory); } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethod.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethod.java index c7806f96..81ae3073 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethod.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethod.java @@ -16,11 +16,13 @@ package org.springframework.graphql.data.method.annotation.support; import java.util.Arrays; +import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executor; import java.util.function.BiConsumer; import graphql.schema.DataFetchingEnvironment; +import io.micrometer.context.ContextSnapshotFactory; import org.reactivestreams.Publisher; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; @@ -45,18 +47,34 @@ public class DataFetcherHandlerMethod extends DataFetcherHandlerMethodSupport { /** - * Constructor with a parent handler method. - * @param handlerMethod the handler method - * @param resolvers the argument resolvers - * @param validationHelper to apply bean validation with - * @param subscription whether the field being fetched is of subscription type + * @deprecated in favor of + * {@link #DataFetcherHandlerMethod(HandlerMethod, HandlerMethodArgumentResolverComposite, BiConsumer, boolean, Executor, ContextSnapshotFactory)} */ + @Deprecated(since = "1.3", forRemoval = true) public DataFetcherHandlerMethod( HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers, @Nullable BiConsumer validationHelper, @Nullable Executor executor, boolean subscription) { - super(handlerMethod, resolvers, executor); + this(handlerMethod, resolvers, validationHelper, subscription, executor, null); + } + + /** + * Create an instance. + * @param handlerMethod the handler method + * @param resolvers the argument resolvers + * @param validationHelper to apply bean validation with + * @param subscription whether the field being fetched is of subscription type + * @param executor {@code Executor} to use for {@link Callable} methods + * @param snapshotFactory for context propagation with {@link Callable} methods + * @since 1.3 + */ + public DataFetcherHandlerMethod( + HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers, + @Nullable BiConsumer validationHelper, boolean subscription, + @Nullable Executor executor, @Nullable ContextSnapshotFactory snapshotFactory) { + + super(handlerMethod, resolvers, executor, snapshotFactory); Assert.isTrue(!resolvers.getResolvers().isEmpty(), "No argument resolvers"); this.validationHelper = (validationHelper != null ? validationHelper : (controller, args) -> {}); this.subscription = subscription; diff --git a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodSupport.java b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodSupport.java index 10be0587..1f5a4251 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodSupport.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodSupport.java @@ -16,9 +16,11 @@ package org.springframework.graphql.data.method.annotation.support; +import java.util.concurrent.Callable; import java.util.concurrent.Executor; import graphql.schema.DataFetchingEnvironment; +import io.micrometer.context.ContextSnapshotFactory; import org.springframework.core.DefaultParameterNameDiscoverer; import org.springframework.core.MethodParameter; @@ -46,11 +48,31 @@ public class DataFetcherHandlerMethodSupport extends InvocableHandlerMethodSuppo private final ParameterNameDiscoverer parameterNameDiscoverer = new DefaultParameterNameDiscoverer(); + /** + * @deprecated in favor of + * {@link #DataFetcherHandlerMethodSupport(HandlerMethod, HandlerMethodArgumentResolverComposite, Executor, ContextSnapshotFactory)} + */ + @Deprecated(since = "1.3", forRemoval = true) protected DataFetcherHandlerMethodSupport( HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers, @Nullable Executor executor) { - super(handlerMethod, executor); + this(handlerMethod, resolvers, executor, null); + } + + /** + * Create an instance. + * @param handlerMethod the handler method + * @param resolvers the argument resolvers + * @param executor {@code Executor} to use for {@link Callable} methods + * @param snapshotFactory for context propagation with {@link Callable} methods + * @since 1.3 + */ + protected DataFetcherHandlerMethodSupport( + HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers, + @Nullable Executor executor, @Nullable ContextSnapshotFactory snapshotFactory) { + + super(handlerMethod, executor, snapshotFactory); this.resolvers = resolvers; } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/AbstractGraphQlSourceBuilder.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/AbstractGraphQlSourceBuilder.java index 535513cb..f2ad217e 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/AbstractGraphQlSourceBuilder.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/AbstractGraphQlSourceBuilder.java @@ -30,6 +30,7 @@ import graphql.schema.GraphQLSchema; import graphql.schema.GraphQLTypeVisitor; import graphql.schema.SchemaTransformer; import graphql.schema.SchemaTraverser; +import io.micrometer.context.ContextSnapshotFactory; import org.springframework.lang.Nullable; @@ -44,6 +45,9 @@ import org.springframework.lang.Nullable; */ public abstract class AbstractGraphQlSourceBuilder> implements GraphQlSource.Builder { + private static final ContextSnapshotFactory DEFAULT_SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); + + private final List exceptionResolvers = new ArrayList<>(); private final List subscriptionExceptionResolvers = new ArrayList<>(); @@ -54,6 +58,9 @@ public abstract class AbstractGraphQlSourceBuilder instrumentations = new ArrayList<>(); + @Nullable + private ContextSnapshotFactory snapshotFactory; + @Nullable private Consumer graphQlConfigurer; @@ -88,6 +95,12 @@ public abstract class AbstractGraphQlSourceBuilder configurer) { this.graphQlConfigurer = (this.graphQlConfigurer != null ? @@ -104,12 +117,15 @@ public abstract class AbstractGraphQlSourceBuilder visitorsToUse = new ArrayList<>(this.typeVisitors); - visitorsToUse.add(ContextDataFetcherDecorator.createVisitor(this.subscriptionExceptionResolvers)); + visitorsToUse.add(contextDataFetcherVisitor); new SchemaTraverser().depthFirstFullSchema(visitorsToUse, schema, vars); return schema.transformWithoutTypes(builder -> builder.codeRegistry(outputCodeRegistry)); 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 a18ec49e..22c737fc 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,19 +59,23 @@ final class ContextDataFetcherDecorator implements DataFetcher { private final SubscriptionExceptionResolver subscriptionExceptionResolver; - private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); + private final ContextSnapshotFactory snapshotFactory; + private ContextDataFetcherDecorator( DataFetcher delegate, boolean subscription, - SubscriptionExceptionResolver subscriptionExceptionResolver) { + SubscriptionExceptionResolver subscriptionExceptionResolver, + ContextSnapshotFactory snapshotFactory) { Assert.notNull(delegate, "'delegate' DataFetcher is required"); Assert.notNull(subscriptionExceptionResolver, "'subscriptionExceptionResolver' is required"); this.delegate = delegate; this.subscription = subscription; this.subscriptionExceptionResolver = subscriptionExceptionResolver; + this.snapshotFactory = snapshotFactory; } + @Override public Object get(DataFetchingEnvironment environment) throws Exception { @@ -113,8 +117,16 @@ final class ContextDataFetcherDecorator implements DataFetcher { * Static factory method to create {@link GraphQLTypeVisitor} that wraps * data fetchers with the {@link ContextDataFetcherDecorator}. */ - static GraphQLTypeVisitor createVisitor(List resolvers) { - return new ContextTypeVisitor(resolvers); + static GraphQLTypeVisitor createVisitor( + List resolvers, ContextSnapshotFactory snapshotFactory) { + + resolvers.forEach(resolver -> { + if (resolver instanceof SubscriptionExceptionResolverAdapter adapter) { + adapter.setContextSnapshotFactory(snapshotFactory); + } + }); + + return new ContextTypeVisitor(resolvers, snapshotFactory); } @@ -125,8 +137,13 @@ final class ContextDataFetcherDecorator implements DataFetcher { private final SubscriptionExceptionResolver exceptionResolver; - private ContextTypeVisitor(List resolvers) { + private final ContextSnapshotFactory snapshotFactory; + + private ContextTypeVisitor( + List resolvers, ContextSnapshotFactory snapshotFactory) { + this.exceptionResolver = new CompositeSubscriptionExceptionResolver(resolvers); + this.snapshotFactory = snapshotFactory; } @Override @@ -142,7 +159,8 @@ final class ContextDataFetcherDecorator implements DataFetcher { if (applyDecorator(dataFetcher)) { boolean handlesSubscription = visitorHelper.isSubscriptionType(parent); - dataFetcher = new ContextDataFetcherDecorator(dataFetcher, handlesSubscription, exceptionResolver); + dataFetcher = new ContextDataFetcherDecorator( + dataFetcher, handlesSubscription, this.exceptionResolver, this.snapshotFactory); codeRegistry.dataFetcher(fieldCoordinates, dataFetcher); } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolver.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolver.java index 724e08ad..4729e5fb 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolver.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolver.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. @@ -22,6 +22,7 @@ import java.util.function.BiFunction; import graphql.GraphQLError; import graphql.execution.DataFetcherExceptionHandler; import graphql.schema.DataFetchingEnvironment; +import io.micrometer.context.ContextSnapshotFactory; import reactor.core.publisher.Mono; /** @@ -101,7 +102,24 @@ public interface DataFetcherExceptionResolver { * @since 1.1.1 */ static DataFetcherExceptionHandler createExceptionHandler(List resolvers) { - return new ExceptionResolversExceptionHandler(resolvers); + return createExceptionHandler(resolvers, ContextSnapshotFactory.builder().build()); + } + + /** + * Variant of {@link #createExceptionHandler(List)} with a + * {@link ContextSnapshotFactory} instance to use. + * @since 1.3 + */ + static DataFetcherExceptionHandler createExceptionHandler( + List resolvers, ContextSnapshotFactory snapshotFactory) { + + resolvers.forEach(resolver -> { + if (resolver instanceof DataFetcherExceptionResolverAdapter adapter) { + adapter.setContextSnapshotFactory(snapshotFactory); + } + }); + + return new ExceptionResolversExceptionHandler(resolvers, snapshotFactory); } } 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 a855432e..064fdffe 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 @@ -50,9 +50,12 @@ import org.springframework.lang.Nullable; */ public abstract class DataFetcherExceptionResolverAdapter implements DataFetcherExceptionResolver { + private static final ContextSnapshotFactory DEFAULT_SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); + + protected final Log logger = LogFactory.getLog(getClass()); - protected final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); + private ContextSnapshotFactory snapshotFactory = DEFAULT_SNAPSHOT_FACTORY; private boolean threadLocalContextAware; @@ -89,6 +92,18 @@ public abstract class DataFetcherExceptionResolverAdapter implements DataFetcher return this.threadLocalContextAware; } + + /** + * Internal method to allow + * via {@link DataFetcherExceptionResolver#createExceptionHandler(List, ContextSnapshotFactory)} + * to set the {@link ContextSnapshotFactory} instance to use. + * @since 1.3 + */ + void setContextSnapshotFactory(ContextSnapshotFactory snapshotFactory) { + this.snapshotFactory = snapshotFactory; + } + + @Override public final Mono> resolveException(Throwable ex, DataFetchingEnvironment env) { return Mono.defer(() -> Mono.justOrEmpty(resolveInternal(ex, env))); @@ -100,7 +115,7 @@ public abstract class DataFetcherExceptionResolverAdapter implements DataFetcher return resolveToMultipleErrors(exception, env); } try { - return snapshotFactory.captureFrom(env.getGraphQlContext()) + return this.snapshotFactory.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 9aabac37..9337a62c 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 @@ -19,6 +19,7 @@ import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionStage; import java.util.function.BiFunction; @@ -53,14 +54,13 @@ 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<>(); private final Supplier defaultOptionsSupplier; + private ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); /** @@ -81,6 +81,26 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry { } + /** + * Configure the {@link ContextSnapshotFactory} instance to use to establish + * {@code ThreadLocal} context for batch loader methods that return {@link Callable}. + * If not set, then an instance with default settings is used. + *
    + *
  • {@link DefaultExecutionGraphQlService#setContextSnapshotFactory} + *
  • {@link GraphQlSource.Builder#contextSnapshotFactory} + *
  • {@link org.springframework.graphql.server.WebGraphQlHandler.Builder#contextSnapshotFactory} + *
+ * @since 1.3 + */ + public void setContextSnapshotFactory(ContextSnapshotFactory snapshotFactory) { + this.snapshotFactory = snapshotFactory; + } + + public ContextSnapshotFactory getContextSnapshotFactory() { + return this.snapshotFactory; + } + + @Override public RegistrationSpec forTypePair(Class keyType, Class valueType) { return new DefaultRegistrationSpec<>(valueType); @@ -161,13 +181,13 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry { @Override public void registerBatchLoader(BiFunction, BatchLoaderEnvironment, Flux> loader) { DefaultBatchLoaderRegistry.this.loaders.add( - new ReactorBatchLoader<>(initName(), loader, initOptionsSupplier())); + new ReactorBatchLoader<>(initName(), loader, initOptionsSupplier(), snapshotFactory)); } @Override public void registerMappedBatchLoader(BiFunction, BatchLoaderEnvironment, Mono>> loader) { DefaultBatchLoaderRegistry.this.mappedLoaders.add( - new ReactorMappedBatchLoader<>(initName(), loader, initOptionsSupplier())); + new ReactorMappedBatchLoader<>(initName(), loader, initOptionsSupplier(), snapshotFactory)); } private String initName() { @@ -209,13 +229,16 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry { private final Supplier optionsSupplier; + private final ContextSnapshotFactory snapshotFactory; + private ReactorBatchLoader(String name, BiFunction, BatchLoaderEnvironment, Flux> loader, - Supplier optionsSupplier) { + Supplier optionsSupplier, ContextSnapshotFactory snapshotFactory) { this.name = name; this.loader = loader; this.optionsSupplier = optionsSupplier; + this.snapshotFactory = snapshotFactory; } public String getName() { @@ -229,7 +252,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 = this.snapshotFactory.captureFrom(graphQLContext); try { return snapshot.wrap(() -> this.loader.apply(keys, environment) @@ -257,13 +280,16 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry { private final Supplier optionsSupplier; + private final ContextSnapshotFactory snapshotFactory; + private ReactorMappedBatchLoader(String name, BiFunction, BatchLoaderEnvironment, Mono>> loader, - Supplier optionsSupplier) { + Supplier optionsSupplier, ContextSnapshotFactory snapshotFactory) { this.name = name; this.loader = loader; this.optionsSupplier = optionsSupplier; + this.snapshotFactory = snapshotFactory; } public String getName() { @@ -277,7 +303,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 = this.snapshotFactory.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 82d1d39f..636900f7 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,12 +46,13 @@ 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; private final List dataLoaderRegistrars = new ArrayList<>(); + private ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); + private boolean hasDataLoaderRegistrations; private final boolean isDefaultExecutionIdProvider; @@ -80,6 +81,23 @@ public class DefaultExecutionGraphQlService implements ExecutionGraphQlService { return !registry.getDataLoaders().isEmpty(); } + /** + * Configure the {@link ContextSnapshotFactory} instance to use to propagate + * {@code TreadLocal} and Reactor context through {@link GraphQLContext}. + * If not set, then an instance with default settings is used. + *

Note that there are other components that would also need to be + * configured similarly to use a single instance: + *

    + *
  • {@link GraphQlSource.Builder#contextSnapshotFactory} + *
  • {@link DefaultBatchLoaderRegistry#setContextSnapshotFactory} + *
  • {@link org.springframework.graphql.server.WebGraphQlHandler.Builder#contextSnapshotFactory} + *
+ * @since 1.3 + */ + public void setContextSnapshotFactory(ContextSnapshotFactory snapshotFactory) { + this.snapshotFactory = snapshotFactory; + } + @Override public final Mono execute(ExecutionGraphQlRequest request) { @@ -91,7 +109,7 @@ public class DefaultExecutionGraphQlService implements ExecutionGraphQlService { ExecutionInput executionInput = request.toExecutionInput(); GraphQLContext graphQLContext = executionInput.getGraphQLContext(); - snapshotFactory.captureFrom(contextView).updateContext(executionInput.getGraphQLContext()); + this.snapshotFactory.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 b0af67fc..20991987 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 @@ -48,17 +48,23 @@ class ExceptionResolversExceptionHandler implements DataFetcherExceptionHandler private static final Log logger = LogFactory.getLog(ExceptionResolversExceptionHandler.class); - private final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); + private final ContextSnapshotFactory snapshotFactory; private final List resolvers; + /** * Create an instance. * @param resolvers the resolvers to use + * @param snapshotFactory the factory instance to use for context propagation */ - ExceptionResolversExceptionHandler(List resolvers) { + ExceptionResolversExceptionHandler( + List resolvers, ContextSnapshotFactory snapshotFactory) { + Assert.notNull(resolvers, "'resolvers' is required"); + Assert.notNull(resolvers, "ContextSnapshotFactory is required"); this.resolvers = new ArrayList<>(resolvers); + this.snapshotFactory = snapshotFactory; } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/GraphQlSource.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/GraphQlSource.java index 119893b0..5495554b 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/execution/GraphQlSource.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/GraphQlSource.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2023 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. @@ -21,12 +21,14 @@ import java.util.function.BiFunction; import java.util.function.Consumer; import graphql.GraphQL; +import graphql.GraphQLContext; import graphql.execution.instrumentation.Instrumentation; import graphql.schema.GraphQLSchema; import graphql.schema.GraphQLTypeVisitor; import graphql.schema.TypeResolver; import graphql.schema.idl.RuntimeWiring; import graphql.schema.idl.TypeDefinitionRegistry; +import io.micrometer.context.ContextSnapshotFactory; import org.springframework.core.io.Resource; @@ -139,6 +141,22 @@ public interface GraphQlSource { */ B instrumentation(List instrumentations); + /** + * Configure the {@link ContextSnapshotFactory} instance to use for + * context propagation of {@code ThreadLocal}, and Reactor context + * values through the {@link GraphQLContext}. + *

Note that there are other components that would also need to be + * configured similarly to use a single instance: + *

    + *
  • {@link DefaultExecutionGraphQlService#setContextSnapshotFactory} + *
  • {@link DefaultBatchLoaderRegistry#setContextSnapshotFactory} + *
  • {@link org.springframework.graphql.server.WebGraphQlHandler.Builder#contextSnapshotFactory} + *
+ * If not set, then an instance with default settings is used. + * @since 1.3 + */ + B contextSnapshotFactory(ContextSnapshotFactory snapshotFactory); + /** * Configure consumers to be given access to the {@link GraphQL.Builder} * used to build {@link GraphQL}. 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 03d4f9ac..012d872e 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 @@ -51,7 +51,7 @@ public abstract class SubscriptionExceptionResolverAdapter implements Subscripti protected final Log logger = LogFactory.getLog(getClass()); - protected final ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); + private ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); private boolean threadLocalContextAware; @@ -80,13 +80,23 @@ public abstract class SubscriptionExceptionResolverAdapter implements Subscripti return this.threadLocalContextAware; } + /** + * Internal method to allow + * via {@link ContextDataFetcherDecorator#createVisitor(List, ContextSnapshotFactory)} + * to set the {@link ContextSnapshotFactory} instance to use. + * @since 1.3 + */ + void setContextSnapshotFactory(ContextSnapshotFactory snapshotFactory) { + this.snapshotFactory = snapshotFactory; + } + @SuppressWarnings({"unused", "try"}) @Override public final Mono> resolveException(Throwable exception) { if (this.threadLocalContextAware) { return Mono.deferContextual(contextView -> { - ContextSnapshot snapshot = snapshotFactory.captureFrom(contextView); + ContextSnapshot snapshot = this.snapshotFactory.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 d5107ce6..049498fd 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 @@ -41,6 +41,9 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { private final List interceptors = new ArrayList<>(); + @Nullable + private ContextSnapshotFactory snapshotFactory; + @Nullable private WebSocketGraphQlInterceptor webSocketInterceptor; @@ -69,9 +72,13 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { } @Override - public WebGraphQlHandler build() { + public WebGraphQlHandler.Builder contextSnapshotFactory(ContextSnapshotFactory snapshotFactory) { + this.snapshotFactory = snapshotFactory; + return this; + } - ContextSnapshotFactory snapshotFactory = ContextSnapshotFactory.builder().build(); + @Override + public WebGraphQlHandler build() { Chain endOfChain = request -> this.service.execute(request).map(WebGraphQlResponse::new); @@ -80,6 +87,9 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { .map(interceptor -> interceptor.apply(endOfChain)) .orElse(endOfChain); + ContextSnapshotFactory snapshotFactory = + (this.snapshotFactory != null ? this.snapshotFactory : ContextSnapshotFactory.builder().build()); + return new WebGraphQlHandler() { @Override @@ -88,6 +98,11 @@ class DefaultWebGraphQlHandlerBuilder implements WebGraphQlHandler.Builder { webSocketInterceptor : new WebSocketGraphQlInterceptor() {}); } + @Override + public ContextSnapshotFactory contextSnapshotFactory() { + return snapshotFactory; + } + @Override public Mono handleRequest(WebGraphQlRequest request) { ContextSnapshot snapshot = snapshotFactory.captureAll(); 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..8430c8e7 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,9 +18,13 @@ package org.springframework.graphql.server; import java.util.List; +import io.micrometer.context.ContextSnapshotFactory; import reactor.core.publisher.Mono; import org.springframework.graphql.ExecutionGraphQlService; +import org.springframework.graphql.execution.DefaultBatchLoaderRegistry; +import org.springframework.graphql.execution.DefaultExecutionGraphQlService; +import org.springframework.graphql.execution.GraphQlSource; /** @@ -39,6 +43,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 +97,22 @@ 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 {@link DefaultExecutionGraphQlService}. + *

Note that there are other components that would also need to be + * configured similarly to use a single instance: + *

    + *
  • {@link DefaultExecutionGraphQlService#setContextSnapshotFactory} + *
  • {@link DefaultBatchLoaderRegistry#setContextSnapshotFactory} + *
  • {@link GraphQlSource.Builder#contextSnapshotFactory} + *
+ * If not set, then an instance with default settings is used. + * @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 88cc3e40..0155ccc0 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 @@ -120,7 +120,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; @@ -358,14 +358,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; + + public ContextHandshakeInterceptor(ContextSnapshotFactory snapshotFactory) { + this.snapshotFactory = snapshotFactory; + } @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/data/method/annotation/support/ContextValueMethodArgumentResolverTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/ContextValueMethodArgumentResolverTests.java index c56efcd0..9d76d2f9 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/ContextValueMethodArgumentResolverTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/ContextValueMethodArgumentResolverTests.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. @@ -124,7 +124,7 @@ public class ContextValueMethodArgumentResolverTests { DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod( new HandlerMethod(new TestController(), TestController.class.getMethod("handleMono", Mono.class)), - resolvers, null, null, false); + resolvers, null, false, null, null); GraphQLContext graphQLContext = new GraphQLContext.Builder().build(); @@ -146,7 +146,7 @@ public class ContextValueMethodArgumentResolverTests { BatchLoaderHandlerMethod handlerMethod = new BatchLoaderHandlerMethod( new HandlerMethod(controller, - TestController.class.getMethod("getAuthors", List.class, Long.class)), null); + TestController.class.getMethod("getAuthors", List.class, Long.class)), null, null); GraphQLContext context = new GraphQLContext.Builder().build(); context.put("id", 123L); diff --git a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java index d4e06b7d..befe4d78 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethodTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2023 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. @@ -59,7 +59,7 @@ public class DataFetcherHandlerMethodTests { resolvers.addResolver(new ArgumentMethodArgumentResolver(new GraphQlArgumentBinder())); DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod( - handlerMethodFor(new TestController(), "hello"), resolvers, null, null, false); + handlerMethodFor(new TestController(), "hello"), resolvers, null, false, null, null); Object result = handlerMethod.invoke( DataFetchingEnvironmentImpl.newDataFetchingEnvironment() @@ -77,7 +77,7 @@ public class DataFetcherHandlerMethodTests { DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod( handlerMethodFor(new TestController(), "handleAndReturnCallable"), resolvers, null, - new SimpleAsyncTaskExecutor(), false); + false, new SimpleAsyncTaskExecutor(), null); DataFetchingEnvironment environment = DataFetchingEnvironmentImpl .newDataFetchingEnvironment() @@ -99,7 +99,7 @@ public class DataFetcherHandlerMethodTests { DataFetcherHandlerMethod handlerMethod = new DataFetcherHandlerMethod( handlerMethodFor(new TestController(), "handleAndReturnFuture"), resolvers, - null, null, false); + null, false, null, null); SecurityContextHolder.setContext(new SecurityContextImpl(new TestingAuthenticationToken("usr", "pwd"))); try {