Refactoring in ReactorContextManager

Use GraphQLContext as input instead of ExecutionInput,
DataFetchingEnvironment and BatchLoaderEnvironment.

See gh-316

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/execution/ContextDataFetcherDecorator.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2022 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 @@ private ContextDataFetcherDecorator(DataFetcher<?> delegate, boolean subscriptio
 
 	@Override
 	public Object get(DataFetchingEnvironment environment) throws Exception {
-		ContextView contextView = ReactorContextManager.getReactorContext(environment);
+		ContextView contextView = ReactorContextManager.getReactorContext(environment.getGraphQlContext());
 
 		Object value;
 		try {

spring-graphql/src/main/java/org/springframework/graphql/execution/DataFetcherExceptionResolverAdapter.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2022 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.
@@ -90,7 +90,7 @@ private List<GraphQLError> resolveInternal(Throwable ex, DataFetchingEnvironment
 		if (!this.threadLocalContextAware) {
 			return resolveToMultipleErrors(ex, env);
 		}
-		ContextView contextView = ReactorContextManager.getReactorContext(env);
+		ContextView contextView = ReactorContextManager.getReactorContext(env.getGraphQlContext());
 		try {
 			ReactorContextManager.restoreThreadLocalValues(contextView);
 			return resolveToMultipleErrors(ex, env);

spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java

@@ -188,7 +188,7 @@ public DataLoaderOptions getOptionsOrDefault(
 
 		@Override
 		public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
-			ContextView contextView = ReactorContextManager.getReactorContext(environment);
+			ContextView contextView = ReactorContextManager.getReactorContext(environment.getContext());
 			try {
 				ReactorContextManager.restoreThreadLocalValues(contextView);
 				return this.loader.apply(keys, environment).collectList().contextWrite(contextView).toFuture();
@@ -239,7 +239,7 @@ public DataLoaderOptions getOptionsOrDefault(
 
 		@Override
 		public CompletionStage<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
-			ContextView contextView = ReactorContextManager.getReactorContext(environment);
+			ContextView contextView = ReactorContextManager.getReactorContext(environment.getContext());
 			try {
 				ReactorContextManager.restoreThreadLocalValues(contextView);
 				return this.loader.apply(keys, environment).contextWrite(contextView).toFuture();

spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultExecutionGraphQlService.java

@@ -76,7 +76,7 @@ public final Mono<ExecutionGraphQlResponse> execute(ExecutionGraphQlRequest requ
 				request.configureExecutionInput(RESET_EXECUTION_ID_CONFIGURER);
 			}
 			ExecutionInput executionInput = request.toExecutionInput();
-			ReactorContextManager.setReactorContext(contextView, executionInput);
+			ReactorContextManager.setReactorContext(contextView, executionInput.getGraphQLContext());
 			ExecutionInput updatedExecutionInput = registerDataLoaders(executionInput);
 			return Mono.fromFuture(this.graphQlSource.graphQl().executeAsync(updatedExecutionInput))
 					.map(result -> new DefaultExecutionGraphQlResponse(updatedExecutionInput, result));

spring-graphql/src/main/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandler.java

@@ -78,7 +78,7 @@ public CompletableFuture<DataFetcherExceptionHandlerResult> handleException(Data
 					.onErrorResume(resolverEx -> Mono.just(handleResolverError(resolverEx, exception, env)))
 					.switchIfEmpty(Mono.fromCallable(() -> createInternalError(exception, env)))
 					.contextWrite((context) -> {
-						ContextView contextView = ReactorContextManager.getReactorContext(env);
+						ContextView contextView = ReactorContextManager.getReactorContext(env.getGraphQlContext());
 						return (contextView.isEmpty() ? context : context.putAll(contextView));
 					})
 					.toFuture();

spring-graphql/src/main/java/org/springframework/graphql/execution/ReactorContextManager.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2022 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,20 +19,17 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 
-import graphql.ExecutionInput;
 import graphql.GraphQLContext;
-import graphql.schema.DataFetchingEnvironment;
-import org.dataloader.BatchLoaderEnvironment;
 import reactor.util.context.Context;
 import reactor.util.context.ContextView;
 
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
 /**
- * Provides helper methods to save Reactor context in the {@link ExecutionInput}
- * so it can be subsequently obtained from {@link DataFetchingEnvironment} and
- * propagated to data fetchers or exception handlers.
+ * Provides helper methods to save Reactor context in the {@link GraphQLContext}
+ * so it can be subsequently obtained and propagated to data fetchers, exception
+ * handlers, and others.
  *
  * <p>The Reactor context is also used to carry ThreadLocal values that are also
  * restored around the execution of data fetchers and exceptions handlers.
@@ -51,36 +48,24 @@ public abstract class ReactorContextManager {
 	private static final String THREAD_LOCAL_ACCESSOR_KEY = ReactorContextManager.class.getName() + ".THREAD_LOCAL_ACCESSOR";
 
 	/**
-	 * Save the given Reactor {@link ContextView} in the an {@link ExecutionInput} for
-	 * later access through the {@link DataFetchingEnvironment}.
-	 * @param contextView the reactor context view
-	 * @param input the input prepared from the GraphQL request
+	 * Save the given Reactor {@link ContextView} in the given {@link GraphQLContext}.
+	 * @param contextView the reactor {@code ContextView} to save
+	 * @param graphQLContext the {@code GraphQLContext} where to save
 	 */
-	static void setReactorContext(ContextView contextView, ExecutionInput input) {
-		input.getGraphQLContext().put(CONTEXT_VIEW_KEY, contextView);
+	static void setReactorContext(ContextView contextView, GraphQLContext graphQLContext) {
+		graphQLContext.put(CONTEXT_VIEW_KEY, contextView);
 	}
 
 	/**
-	 * Return the Reactor {@link ContextView} saved in the given DataFetchingEnvironment.
-	 * @param environment the DataFetchingEnvironment
+	 * Return the Reactor {@link ContextView} saved in the given {@link GraphQLContext}.
+	 * @param graphQlContext the DataFetchingEnvironment
 	 * @return the reactor {@link ContextView}
 	 */
-	static ContextView getReactorContext(DataFetchingEnvironment environment) {
-		GraphQLContext graphQlContext = environment.getGraphQlContext();
+	static ContextView getReactorContext(GraphQLContext graphQlContext) {
+		Assert.notNull(graphQlContext, "GraphQLContext is required");
 		return graphQlContext.getOrDefault(CONTEXT_VIEW_KEY, Context.empty());
 	}
 
-	/**
-	 * Return the Reactor {@link ContextView} saved in the given BatchLoaderEnvironment.
-	 * @param environment the BatchLoaderEnvironment
-	 * @return the reactor {@link ContextView}
-	 */
-	static ContextView getReactorContext(BatchLoaderEnvironment environment) {
-		Object context = environment.getContext();
-		Assert.isTrue(context instanceof GraphQLContext, "Expected GraphQLContext in BatchLoaderEnvironment");
-		return ((GraphQLContext) context).getOrDefault(CONTEXT_VIEW_KEY, Context.empty());
-	}
-
 	/**
 	 * Use the given accessor to extract ThreadLocal values and save them in a
 	 * sub-map in the given {@link Context}, so those can be restored later

spring-graphql/src/test/java/org/springframework/graphql/execution/ContextDataFetcherDecoratorTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2022 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.
@@ -52,7 +52,7 @@ void monoDataFetcher() throws Exception {
 				.toGraphQl();
 
 		ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
-		ReactorContextManager.setReactorContext(Context.of("name", "007"), input);
+		ReactorContextManager.setReactorContext(Context.of("name", "007"), input.getGraphQLContext());
 
 		ExecutionResult executionResult = graphQl.executeAsync(input).get();
 
@@ -72,7 +72,7 @@ void fluxDataFetcher() throws Exception {
 				.toGraphQl();
 
 		ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greetings }").build();
-		ReactorContextManager.setReactorContext(Context.of("name", "007"), input);
+		ReactorContextManager.setReactorContext(Context.of("name", "007"), input.getGraphQLContext());
 
 		ExecutionResult result = graphQl.executeAsync(input).get();
 
@@ -92,7 +92,7 @@ void fluxDataFetcherSubscription() throws Exception {
 				.toGraphQl();
 
 		ExecutionInput input = ExecutionInput.newExecutionInput().query("subscription { greetings }").build();
-		ReactorContextManager.setReactorContext(Context.of("name", "007"), input);
+		ReactorContextManager.setReactorContext(Context.of("name", "007"), input.getGraphQLContext());
 
 		ExecutionResult executionResult = graphQl.executeAsync(input).get();
 
@@ -116,7 +116,7 @@ void dataFetcherWithThreadLocalContext() {
 
 			ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
 			ContextView view = ReactorContextManager.extractThreadLocalValues(accessor, Context.empty());
-			ReactorContextManager.setReactorContext(view, input);
+			ReactorContextManager.setReactorContext(view, input.getGraphQLContext());
 
 			Mono<ExecutionResult> resultMono = Mono.delay(Duration.ofMillis(10))
 					.flatMap((aLong) -> Mono.fromFuture(graphQl.executeAsync(input)));

spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2021 the original author or authors.
+ * Copyright 2002-2022 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.
@@ -111,7 +111,7 @@ void batchLoaderWithCustomNameAndOptions() {
 
 	private GraphQLContext initGraphQLContext(ContextView context) {
 		ExecutionInput executionInput = ExecutionInput.newExecutionInput().query("").build();
-		ReactorContextManager.setReactorContext(context, executionInput);
+		ReactorContextManager.setReactorContext(context, executionInput.getGraphQLContext());
 		return executionInput.getGraphQLContext();
 	}
 

spring-graphql/src/test/java/org/springframework/graphql/execution/ExceptionResolversExceptionHandlerTests.java

@@ -76,7 +76,7 @@ void resolveExceptionWithReactorContext() throws Exception {
 								.message("Resolved error: " + ex.getMessage() + ", name=" + view.get("name"))
 								.errorType(ErrorType.BAD_REQUEST).build())));
 
-		ReactorContextManager.setReactorContext(Context.of("name", "007"), input);
+		ReactorContextManager.setReactorContext(Context.of("name", "007"), input.getGraphQLContext());
 
 		ExecutionResult result = this.graphQlSetup.exceptionResolver(resolver).toGraphQl()
 				.executeAsync(this.input).get();
@@ -102,7 +102,7 @@ void resolveExceptionWithThreadLocal() {
 			resolver.setThreadLocalContextAware(true);
 
 			ContextView view = ReactorContextManager.extractThreadLocalValues(accessor, Context.empty());
-			ReactorContextManager.setReactorContext(view, input);
+			ReactorContextManager.setReactorContext(view, input.getGraphQLContext());
 
 			Mono<ExecutionResult> result = Mono.delay(Duration.ofMillis(10)).flatMap((aLong) ->
 					Mono.fromFuture(this.graphQlSetup.exceptionResolver(resolver).toGraphQl().executeAsync(this.input)));