Support Principal as a method argument

Closes gh-119

spring-graphql-docs/src/docs/asciidoc/index.adoc

@@ -605,6 +605,9 @@ See <<controllers-schema-mapping-data-loader>>.
 | `GraphQLContext`
 | For access to the context from the `DataFetchingEnvironment`.
 
+| `java.security.Principal`
+| Obtained from Spring Security context, if available.
+
 | `DataFetchingFieldSelectionSet`
 | For access to the selection set for the query through the `DataFetchingEnvironment`.
 
@@ -802,6 +805,9 @@ Batch mapping methods support two types of arguments:
 | `List<K>`
 | The source/parent objects.
 
+| `java.security.Principal`
+| Obtained from Spring Security context, if available.
+
 | `BatchLoaderEnvironment`
 | The environment that is available in GraphQL Java to a
 `org.dataloader.BatchLoaderWithContext`.

spring-graphql/build.gradle

@@ -32,6 +32,7 @@ dependencies {
 	testImplementation 'org.springframework:spring-test'
 	testImplementation 'org.springframework.data:spring-data-commons'
 	testImplementation 'org.springframework.data:spring-data-keyvalue'
+	testImplementation 'org.springframework.security:spring-security-core'
 	testImplementation 'com.querydsl:querydsl-core'
 	testImplementation 'com.querydsl:querydsl-collections'
 	testImplementation 'javax.servlet:javax.servlet-api'

spring-graphql/src/main/java/org/springframework/graphql/data/method/HandlerMethodArgumentResolver.java

@@ -19,13 +19,14 @@ import graphql.schema.DataFetchingEnvironment;
 
 import org.springframework.core.MethodParameter;
 import org.springframework.lang.Nullable;
-import org.springframework.web.bind.WebDataBinder;
-import org.springframework.web.bind.support.WebDataBinderFactory;
-import org.springframework.web.method.support.ModelAndViewContainer;
 
 /**
  * Strategy interface for resolving method parameters into argument values in
- * the context of a given request.
+ * the context of a given {@link DataFetchingEnvironment}.
+ *
+ * <p>Most implementations will be synchronous, simply resolving values from the
+ * {@code DataFetchingEnvironment}. However, a resolver may also return a
+ * {@link reactor.core.publisher.Mono} if it needs to be asynchronous.
  *
  * @author Rossen Stoyanchev
  * @since 1.0.0
@@ -33,25 +34,21 @@ import org.springframework.web.method.support.ModelAndViewContainer;
 public interface HandlerMethodArgumentResolver {
 
 	/**
-	 * Whether the given {@linkplain MethodParameter method parameter} is
-	 * supported by this resolver.
-	 * @param parameter the method parameter to check
-	 * @return {@code true} if this resolver supports the supplied parameter;
-	 * {@code false} otherwise
+	 * Whether this resolver supports the given {@link MethodParameter}.
 	 */
 	boolean supportsParameter(MethodParameter parameter);
 
 	/**
-	 * Resolves a method parameter into an argument value from a given request.
-	 * A {@link ModelAndViewContainer} provides access to the model for the
-	 * request. A {@link WebDataBinderFactory} provides a way to create
-	 * a {@link WebDataBinder} instance when needed for data binding and
-	 * type conversion purposes.
+	 * Resolve a method parameter to a value.
+	 *
 	 * @param parameter the method parameter to resolve. This parameter must
-	 * have previously been passed to {@link #supportsParameter} which must
-	 * have returned {@code true}.
-	 * @param environment the GraphQL {@link DataFetchingEnvironment}
-	 * @return the resolved argument value, or {@code null} if not resolvable
+	 * have previously checked via {@link #supportsParameter}.
+	 * @param environment the environment to use to resolve the value
+	 *
+	 * @return the resolved value, which may be {@code null} if not resolved;
+	 * the value may also be a {@link reactor.core.publisher.Mono} if it
+	 * requires asynchronous resolution.
+	 *
 	 * @throws Exception in case of errors with the preparation of argument values
 	 */
 	@Nullable

spring-graphql/src/main/java/org/springframework/graphql/data/method/InvocableHandlerMethodSupport.java

@@ -17,20 +17,31 @@ package org.springframework.graphql.data.method;
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import org.springframework.core.CoroutinesUtils;
 import org.springframework.core.KotlinDetector;
 import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
 
 /**
  * Extension of {@link HandlerMethod} that adds support for invoking the
- * annotated handler methods.
+ * underlying handler methods.
  *
  * @author Rossen Stoyanchev
  * @since 1.0.0
  */
 public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
 
+	private static final Object NO_VALUE = new Object();
+
 
 	protected InvocableHandlerMethodSupport(HandlerMethod handlerMethod) {
 		super(handlerMethod.createWithResolvedBean());
@@ -39,37 +50,59 @@ public abstract class InvocableHandlerMethodSupport extends HandlerMethod {
 
 	/**
 	 * Invoke the handler method with the given argument values.
+	 * @param argValues the values to use to invoke the method
+	 * @return the value returned from the method or a {@code Mono<Throwable>}
+	 * if the invocation fails.
 	 */
 	@Nullable
-	protected Object doInvoke(Object... args) throws Exception {
+	protected Object doInvoke(Object... argValues) {
+		if (logger.isTraceEnabled()) {
+			logger.trace("Arguments: " + Arrays.toString(argValues));
+		}
 		Method method = getBridgedMethod();
 		try {
 			if (KotlinDetector.isSuspendingFunction(method)) {
-				return CoroutinesUtils.invokeSuspendingFunction(method, getBean(), args);
+				return CoroutinesUtils.invokeSuspendingFunction(method, getBean(), argValues);
 			}
-			return method.invoke(getBean(), args);
+			return method.invoke(getBean(), argValues);
 		}
 		catch (IllegalArgumentException ex) {
-			assertTargetBean(method, getBean(), args);
+			assertTargetBean(method, getBean(), argValues);
 			String text = (ex.getMessage() != null ? ex.getMessage() : "Illegal argument");
-			throw new IllegalStateException(formatInvokeError(text, args), ex);
+			return Mono.error(new IllegalStateException(formatInvokeError(text, argValues), ex));
 		}
 		catch (InvocationTargetException ex) {
 			// Unwrap for DataFetcherExceptionResolvers ...
 			Throwable targetException = ex.getTargetException();
-			if (targetException instanceof RuntimeException) {
-				throw (RuntimeException) targetException;
-			}
-			else if (targetException instanceof Error) {
-				throw (Error) targetException;
-			}
-			else if (targetException instanceof Exception) {
-				throw (Exception) targetException;
+			if (targetException instanceof Error || targetException instanceof Exception) {
+				return Mono.error(targetException);
 			}
 			else {
-				throw new IllegalStateException(formatInvokeError("Invocation failure", args), targetException);
+				return Mono.error(new IllegalStateException(
+						formatInvokeError("Invocation failure", argValues), targetException));
 			}
 		}
+		catch (Throwable ex) {
+			return Mono.error(ex);
+		}
+	}
+
+	/**
+	 * Use this method to resolve the arguments asynchronously. This is only
+	 * useful when at least one of the values is a {@link Mono}
+	 */
+	@SuppressWarnings("unchecked")
+	protected Mono<Object[]> toArgsMono(Object[] args) {
+
+		List<Mono<Object>> monoList = Arrays.stream(args)
+				.map(arg -> {
+					Mono<Object> argMono = (arg instanceof Mono ? (Mono<Object>) arg : Mono.just(arg));
+					return argMono.defaultIfEmpty(NO_VALUE);
+				})
+				.collect(Collectors.toList());
+
+		return Mono.zip(monoList,
+				values -> Stream.of(values).map(value -> value != NO_VALUE ? value : null).toArray());
 	}
 
 }

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AnnotatedControllerConfigurer.java

@@ -84,6 +84,10 @@ public class AnnotatedControllerConfigurer
 	 */
 	private static final String SCOPED_TARGET_NAME_PREFIX = "scopedTarget.";
 
+	private final static boolean springSecurityPresent = ClassUtils.isPresent(
+			"org.springframework.security.core.context.SecurityContext",
+			AnnotatedControllerConfigurer.class.getClassLoader());
+
 
 	@Nullable
 	private ApplicationContext applicationContext;
@@ -117,6 +121,9 @@ public class AnnotatedControllerConfigurer
 		this.argumentResolvers.addResolver(new ArgumentMethodArgumentResolver(this.conversionService));
 		this.argumentResolvers.addResolver(new DataFetchingEnvironmentMethodArgumentResolver());
 		this.argumentResolvers.addResolver(new DataLoaderMethodArgumentResolver());
+		if (springSecurityPresent) {
+			this.argumentResolvers.addResolver(new PrincipalMethodArgumentResolver());
+		}
 
 		if (KotlinDetector.isKotlinPresent()) {
 			this.argumentResolvers.addResolver(new ContinuationHandlerMethodArgumentResolver());
@@ -288,28 +295,27 @@ public class AnnotatedControllerConfigurer
 				.collect(Collectors.joining("\n\t", "\n\t" + formattedType + ":" + "\n\t", ""));
 	}
 
-	@SuppressWarnings("unchecked")
-	private <P, F> String registerBatchLoader(MappingInfo info) {
+	private String registerBatchLoader(MappingInfo info) {
 		if (!info.isBatchMapping()) {
 			throw new IllegalArgumentException("Not a @BatchMapping method: " + info);
 		}
 
 		String dataLoaderKey = info.getCoordinates().toString();
-		BatchLoaderHandlerMethod invocable = new BatchLoaderHandlerMethod(info.getHandlerMethod());
 		BatchLoaderRegistry registry = obtainApplicationContext().getBean(BatchLoaderRegistry.class);
 
-		Class<?> clazz = info.getHandlerMethod().getReturnType().getParameterType();
+		HandlerMethod handlerMethod = info.getHandlerMethod();
+		BatchLoaderHandlerMethod invocable = new BatchLoaderHandlerMethod(handlerMethod);
+
+		Class<?> clazz = handlerMethod.getReturnType().getParameterType();
 		if (clazz.equals(Flux.class) || Collection.class.isAssignableFrom(clazz)) {
-			registry.<P,F>forName(dataLoaderKey).registerBatchLoader((values, env) ->
-					(Flux<F>) invocable.invoke(values, env));
+			registry.forName(dataLoaderKey).registerBatchLoader(invocable::invokeForIterable);
 		}
 		else if (clazz.equals(Mono.class) || clazz.equals(Map.class)) {
-			registry.<P,F>forName(dataLoaderKey).registerMappedBatchLoader((values, env) ->
-					(Mono<Map<P, F>>) invocable.invoke(values, env));
+			registry.forName(dataLoaderKey).registerMappedBatchLoader(invocable::invokeForMap);
 		}
 		else {
 			throw new IllegalStateException("@BatchMapping method is expected to return " +
-					"Flux<V>, List<V>, Mono<Map<K, V>>, or Map<K, V>: " + info.getHandlerMethod());
+					"Flux<V>, List<V>, Mono<Map<K, V>>, or Map<K, V>: " + handlerMethod);
 		}
 
 		return dataLoaderKey;
@@ -358,9 +364,12 @@ public class AnnotatedControllerConfigurer
 
 		private final HandlerMethodArgumentResolverComposite argumentResolvers;
 
+		private final boolean subscription;
+
 		public SchemaMappingDataFetcher(MappingInfo info, HandlerMethodArgumentResolverComposite resolvers) {
 			this.info = info;
 			this.argumentResolvers = resolvers;
+			this.subscription = this.info.getCoordinates().getTypeName().equalsIgnoreCase("Subscription");
 		}
 
 		/**
@@ -381,7 +390,7 @@ public class AnnotatedControllerConfigurer
 		@Override
 		@SuppressWarnings("ConstantConditions")
 		public Object get(DataFetchingEnvironment environment) throws Exception {
-			return new DataFetcherHandlerMethod(getHandlerMethod(), this.argumentResolvers).invoke(environment);
+			return new DataFetcherHandlerMethod(getHandlerMethod(), this.argumentResolvers, this.subscription).invoke(environment);
 		}
 	}
 

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java

@@ -15,6 +15,8 @@
  */
 package org.springframework.graphql.data.method.annotation.support;
 
+import java.security.Principal;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Map;
 
@@ -27,7 +29,7 @@ import org.springframework.core.MethodParameter;
 import org.springframework.graphql.data.method.HandlerMethod;
 import org.springframework.graphql.data.method.InvocableHandlerMethodSupport;
 import org.springframework.lang.Nullable;
-import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
 
 /**
  * An extension of {@link HandlerMethod} for annotated handler methods adapted to
@@ -40,6 +42,10 @@ import org.springframework.util.Assert;
  */
 public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 
+	private final static boolean springSecurityPresent = ClassUtils.isPresent(
+			"org.springframework.security.core.context.SecurityContext",
+			AnnotatedControllerConfigurer.class.getClassLoader());
+
 
 	public BatchLoaderHandlerMethod(HandlerMethod handlerMethod) {
 		super(handlerMethod);
@@ -47,43 +53,55 @@ public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 
 
 	/**
-	 * Invoke the underlying batch loading method, resolving its arguments from
-	 * the given keys  and the {@link BatchLoaderEnvironment}.
+	 * Invoke the underlying batch loader method with a collection of keys to
+	 * return a Map of key-value pairs.
 	 *
-	 * @param keys the batch loading keys
+	 * @param keys the keys for which to load values
 	 * @param environment the environment available to batch loaders
-	 * @return a {@code Flux} of values or a {@code Mono} with map of key-value pairs.
+	 * @param <K> the type of keys in the map
+	 * @param <V> the type of values in the map
+	 * @return a {@code Mono} with map of key-value pairs.
 	 */
 	@Nullable
-	public <K> Object invoke(Collection<K> keys, BatchLoaderEnvironment environment) {
-
-		MethodParameter[] parameters = getMethodParameters();
-		Assert.notEmpty(parameters, "Batch loading methods should have at least " +
-				"one argument with the List of parent objects: " + getBridgedMethod().toGenericString());
-
-		Object[] args = new Object[parameters.length];
-		for (int i = 0; i < parameters.length; i++) {
-			args[i] = resolveArgument(parameters[i], keys, environment);
+	public <K, V> Mono<Map<K, V>> invokeForMap(Collection<K> keys, BatchLoaderEnvironment environment) {
+		Object[] args = getMethodArgumentValues(keys, environment);
+		if (doesNotHaveAsyncArgs(args)) {
+			Object result = doInvoke(args);
+			return toMonoMap(result);
 		}
+		return toArgsMono(args).flatMap(argValues -> {
+			Object result = doInvoke(argValues);
+			return toMonoMap(result);
+		});
+	}
 
-		Object result;
-		try {
-			result = doInvoke(args);
-		}
-		catch (Exception ex) {
-			throw new IllegalStateException("...", ex);
+	/**
+	 * Invoke the underlying batch loader method with a collection of input keys
+	 * to return a collection of matching values.
+	 *
+	 * @param keys the keys for which to load values
+	 * @param environment the environment available to batch loaders
+	 * @param <V> the type of values returned
+	 * @return a {@code Flux} of values.
+	 */
+	public <V> Flux<V> invokeForIterable(Collection<?> keys, BatchLoaderEnvironment environment) {
+		Object[] args = getMethodArgumentValues(keys, environment);
+		if (doesNotHaveAsyncArgs(args)) {
+			Object result = doInvoke(args);
+			return toFlux(result);
 		}
+		return toArgsMono(args).flatMapMany(resolvedArgs -> {
+			Object result = doInvoke(resolvedArgs);
+			return toFlux(result);
+		});
+	}
 
-		if (result != null) {
-			if (result instanceof Collection) {
-				return Flux.fromIterable((Collection<?>) result);
-			}
-			else if (result instanceof Map) {
-				return Mono.just(result);
-			}
+	private <K> Object[] getMethodArgumentValues(Collection<K> keys, BatchLoaderEnvironment environment) {
+		Object[] args = new Object[getMethodParameters().length];
+		for (int i = 0; i < getMethodParameters().length; i++) {
+			args[i] = resolveArgument(getMethodParameters()[i], keys, environment);
 		}
-
-		return result;
+		return args;
 	}
 
 	@Nullable
@@ -107,9 +125,38 @@ public class BatchLoaderHandlerMethod extends InvocableHandlerMethodSupport {
 		else if ("kotlin.coroutines.Continuation".equals(parameterType.getName())) {
 			return null;
 		}
+		else if (springSecurityPresent && Principal.class.isAssignableFrom(parameter.getParameterType())) {
+			return PrincipalMethodArgumentResolver.doResolve();
+		}
 		else {
 			throw new IllegalStateException(formatArgumentError(parameter, "Unexpected argument type."));
 		}
 	}
 
+	private boolean doesNotHaveAsyncArgs(Object[] args) {
+		return Arrays.stream(args).noneMatch(arg -> arg instanceof Mono);
+	}
+
+	@SuppressWarnings("unchecked")
+	private static <K, V> Mono<Map<K, V>> toMonoMap(@Nullable Object result) {
+		if (result instanceof Map) {
+			return Mono.just((Map<K, V>) result);
+		}
+		else if (result instanceof Mono) {
+			return (Mono<Map<K, V>>) result;
+		}
+		return Mono.error(new IllegalStateException("Unexpected return value: " + result));
+	}
+
+	@SuppressWarnings("unchecked")
+	private static <V> Flux<V> toFlux(@Nullable Object result) {
+		if (result instanceof Collection) {
+			return Flux.fromIterable((Collection<V>) result);
+		}
+		else if (result instanceof Flux) {
+			return (Flux<V>) result;
+		}
+		return Flux.error(new IllegalStateException("Unexpected return value: " + result));
+	}
+
 }

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/DataFetcherHandlerMethod.java

@@ -18,6 +18,9 @@ package org.springframework.graphql.data.method.annotation.support;
 import java.util.Arrays;
 
 import graphql.schema.DataFetchingEnvironment;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
 import org.springframework.core.DefaultParameterNameDiscoverer;
 import org.springframework.core.MethodParameter;
@@ -47,11 +50,16 @@ public class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
 
 	private final ParameterNameDiscoverer parameterNameDiscoverer = new DefaultParameterNameDiscoverer();
 
+	private final boolean subscription;
+
+
+	public DataFetcherHandlerMethod(
+			HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers, boolean subscription) {
 
-	public DataFetcherHandlerMethod(HandlerMethod handlerMethod, HandlerMethodArgumentResolverComposite resolvers) {
 		super(handlerMethod);
 		Assert.isTrue(!resolvers.getResolvers().isEmpty(), "No argument resolvers");
 		this.resolvers = resolvers;
+		this.subscription = subscription;
 	}
 
 
@@ -66,25 +74,53 @@ public class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
 	/**
 	 * Invoke the method after resolving its argument values in the context of
 	 * the given {@link DataFetchingEnvironment}.
+	 *
 	 * <p>Argument values are commonly resolved through
 	 * {@link HandlerMethodArgumentResolver HandlerMethodArgumentResolvers}.
 	 * The {@code providedArgs} parameter however may supply argument values to
 	 * be used directly, i.e. without argument resolution. Provided argument
 	 * values are checked before argument resolvers.
-	 * @param environment the GraphQL {@link DataFetchingEnvironment}
-	 * @return the raw value returned by the invoked method
-	 * @throws Exception raised if no suitable argument resolver can be found,
-	 * or if the method raised an exception
-	 * @see #getMethodArgumentValues
-	 * @see #doInvoke
+	 *
+	 * @param environment the GraphQL {@link DataFetchingEnvironment} to use to
+	 * resolve arguments.
+	 *
+	 * @return the raw value returned by the invoked method, which may also be
+	 * wrapped as a {@code Mono} in case of method arguments that require
+	 * asynchronous resolution, e.g. {@code Principal} in WebFlux; this method
+	 * may also return a {@code Mono<Throwable>} if the invocation fails.
 	 */
 	@Nullable
 	public Object invoke(DataFetchingEnvironment environment) throws Exception {
-		Object[] args = getMethodArgumentValues(environment);
-		if (logger.isTraceEnabled()) {
-			logger.trace("Arguments: " + Arrays.toString(args));
+		Object[] args;
+		try {
+			args = getMethodArgumentValues(environment);
 		}
-		return doInvoke(args);
+		catch (Throwable ex) {
+			return Mono.error(ex);
+		}
+
+		if (Arrays.stream(args).noneMatch(arg -> arg instanceof Mono)) {
+			return doInvoke(args);
+		}
+
+		return this.subscription ?
+				toArgsMono(args).flatMapMany(argValues -> {
+					Object result = doInvoke(argValues);
+					Assert.state(result instanceof Publisher, "Expected a Publisher from a Subscription response");
+					return Flux.from((Publisher<?>) result);
+				}) :
+				toArgsMono(args).flatMap(argValues -> {
+					Object result = doInvoke(argValues);
+					if (result instanceof Mono) {
+						return (Mono<?>) result;
+					}
+					else if (result instanceof Flux) {
+						return Flux.from((Flux<?>) result).collectList();
+					}
+					else {
+						return Mono.justOrEmpty(result);
+					}
+				});
 	}
 
 	/**
@@ -92,13 +128,14 @@ public class DataFetcherHandlerMethod extends InvocableHandlerMethodSupport {
 	 * argument values and falling back to the configured argument resolvers.
 	 * <p>The resulting array will be passed into {@link #doInvoke}.
 	 */
-	protected Object[] getMethodArgumentValues(
+	private Object[] getMethodArgumentValues(
 			DataFetchingEnvironment environment, Object... providedArgs) throws Exception {
 
 		MethodParameter[] parameters = getMethodParameters();
 		if (ObjectUtils.isEmpty(parameters)) {
 			return EMPTY_ARGS;
 		}
+
 		Object[] args = new Object[parameters.length];
 		for (int i = 0; i < parameters.length; i++) {
 			MethodParameter parameter = parameters[i];

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/PrincipalMethodArgumentResolver.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2002-2021 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.data.method.annotation.support;
+
+import java.security.Principal;
+
+import graphql.schema.DataFetchingEnvironment;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.ReactiveSecurityContextHolder;
+import org.springframework.security.core.context.SecurityContext;
+import org.springframework.security.core.context.SecurityContextHolder;
+
+/**
+ * Resolver to obtain {@link Principal} from Spring Security context via
+ * {@link SecurityContext#getAuthentication()}.
+ *
+ * <p>The resolver checks both ThreadLocal context via {@link SecurityContextHolder}
+ * for Spring MVC applications, and {@link ReactiveSecurityContextHolder} for
+ * Spring WebFlux applications. It returns .
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.0.0
+ */
+public class PrincipalMethodArgumentResolver implements HandlerMethodArgumentResolver {
+
+	/**
+	 * Return "true" if the argument is {@link Principal} or a sub-type.
+	 */
+	@Override
+	public boolean supportsParameter(MethodParameter parameter) {
+		return Principal.class.isAssignableFrom(parameter.getParameterType());
+	}
+
+
+	@Override
+	public Object resolveArgument(MethodParameter parameter, DataFetchingEnvironment environment) {
+		return doResolve();
+	}
+
+	static Object doResolve() {
+		Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
+		return (authentication != null ? authentication :
+				ReactiveSecurityContextHolder.getContext().map(SecurityContext::getAuthentication));
+	}
+
+}

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

@@ -179,7 +179,13 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
 		@Override
 		public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment environment) {
 			ContextView contextView = ReactorContextManager.getReactorContext(environment);
-			return this.loader.apply(keys, environment).collectList().contextWrite(contextView).toFuture();
+			try {
+				ReactorContextManager.restoreThreadLocalValues(contextView);
+				return this.loader.apply(keys, environment).collectList().contextWrite(contextView).toFuture();
+			}
+			finally {
+				ReactorContextManager.resetThreadLocalValues(contextView);
+			}
 		}
 
 	}
@@ -217,7 +223,13 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
 		@Override
 		public CompletionStage<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment environment) {
 			ContextView contextView = ReactorContextManager.getReactorContext(environment);
-			return this.loader.apply(keys, environment).contextWrite(contextView).toFuture();
+			try {
+				ReactorContextManager.restoreThreadLocalValues(contextView);
+				return this.loader.apply(keys, environment).contextWrite(contextView).toFuture();
+			}
+			finally {
+				ReactorContextManager.resetThreadLocalValues(contextView);
+			}
 		}
 
 	}

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/BatchMappingPrincipalMethodArgumentResolverTests.java

@@ -0,0 +1,195 @@
+/*
+ * Copyright 2002-2021 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.data.method.annotation.support;
+
+import java.security.Principal;
+import java.time.Duration;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import graphql.ExecutionResult;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.util.context.Context;
+
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.graphql.GraphQlResponse;
+import org.springframework.graphql.RequestInput;
+import org.springframework.graphql.data.method.annotation.BatchMapping;
+import org.springframework.graphql.data.method.annotation.QueryMapping;
+import org.springframework.graphql.execution.ExecutionGraphQlService;
+import org.springframework.graphql.execution.ReactorContextManager;
+import org.springframework.graphql.security.SecurityContextThreadLocalAccessor;
+import org.springframework.lang.Nullable;
+import org.springframework.security.authentication.TestingAuthenticationToken;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.ReactiveSecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextImpl;
+import org.springframework.stereotype.Controller;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Named.named;
+import static org.junit.jupiter.params.provider.Arguments.arguments;
+
+/**
+ * Tests for {@code @BatchMapping} methods with a {@link Principal} argument.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class BatchMappingPrincipalMethodArgumentResolverTests extends BatchMappingTestSupport {
+
+	private final Authentication authentication = new TestingAuthenticationToken(new Object(), new Object());
+
+	private final Function<Context, Context> reactiveContextWriter = context ->
+			ReactiveSecurityContextHolder.withAuthentication(this.authentication);
+
+	private final Function<Context, Context> threadLocalContextWriter = context ->
+			ReactorContextManager.extractThreadLocalValues(new SecurityContextThreadLocalAccessor(), context);
+
+
+	private static Stream<Arguments> controllers() {
+		return Stream.of(
+				arguments(named("Returning Mono<Map<K,V>>", new BatchMonoMapController())),
+				arguments(named("Returning Map<K,V>", new BatchMapController())),
+				arguments(named("Returning Flux<V>", new BatchFluxController())),
+				arguments(named("Returning List<V>", new BatchListController()))
+		);
+	}
+
+	@ParameterizedTest
+	@MethodSource("controllers")
+	void resolveFromReactiveContext(CourseController courseController) {
+		testBatchLoading(courseController, this.reactiveContextWriter);
+	}
+
+	@ParameterizedTest
+	@MethodSource("controllers")
+	void resolveFromThreadLocalContext(CourseController courseController) {
+		SecurityContextHolder.setContext(new SecurityContextImpl(authentication));
+		try {
+			testBatchLoading(courseController, this.threadLocalContextWriter);
+		}
+		finally {
+			SecurityContextHolder.clearContext();
+		}
+	}
+
+	private void testBatchLoading(CourseController controller, Function<Context, Context> contextWriter) {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+		context.getBeanFactory().registerSingleton("courseController", controller);
+		context.register(BatchMappingTestSupport.CourseConfig.class);
+		context.refresh();
+
+		ExecutionGraphQlService graphQlService = context.getBean(ExecutionGraphQlService.class);
+
+		Mono<ExecutionResult> resultMono = Mono.delay(Duration.ofMillis(10))
+				.flatMap(aLong -> {
+					String query = "{ courses { id instructor { id } } }";
+					return graphQlService.execute(new RequestInput(query, null, null, null));
+				})
+				.contextWrite(contextWriter);
+
+		List<Course> actualCourses = GraphQlResponse.from(resultMono).toList("courses", Course.class);
+		List<Course> courses = Course.allCourses();
+		assertThat(actualCourses).hasSize(courses.size());
+		for (int i = 0; i < courses.size(); i++) {
+			assertThat(actualCourses.get(i).instructor()).isEqualTo(courses.get(i).instructor());
+		}
+
+		assertThat(controller.principal()).isSameAs(this.authentication);
+	}
+
+
+	@SuppressWarnings("unused")
+	private static class CourseController {
+
+		@Nullable
+		protected Principal principal;
+
+		@Nullable
+		public Principal principal() {
+			return this.principal;
+		}
+
+		protected void principal(Principal principal) {
+			this.principal = principal;
+		}
+
+		@QueryMapping
+		public Collection<Course> courses() {
+			return BatchMappingTestSupport.courseMap.values();
+		}
+	}
+
+
+	@Controller
+	@SuppressWarnings("unused")
+	private static class BatchMonoMapController extends CourseController {
+
+		@BatchMapping
+		public Mono<Map<Course, Person>> instructor(List<Course> courses, Principal principal) {
+			principal(principal);
+			return Flux.fromIterable(courses).collect(Collectors.toMap(Function.identity(), Course::instructor));
+		}
+
+	}
+
+
+	@Controller
+	@SuppressWarnings("unused")
+	private static class BatchMapController extends CourseController {
+
+		@BatchMapping
+		public Map<Course, Person> instructor(List<Course> courses, Principal principal) {
+			principal(principal);
+			return courses.stream().collect(Collectors.toMap(Function.identity(), Course::instructor));
+		}
+
+	}
+
+	@Controller
+	@SuppressWarnings("unused")
+	private static class BatchFluxController extends CourseController {
+
+		@BatchMapping
+		public Flux<Person> instructor(List<Course> courses, Principal principal) {
+			principal(principal);
+			return Flux.fromIterable(courses).map(Course::instructor);
+		}
+
+	}
+
+	@Controller
+	@SuppressWarnings("unused")
+	private static class BatchListController extends CourseController {
+
+		@BatchMapping
+		public List<Person> instructor(List<Course> courses, Principal principal) {
+			principal(principal);
+			return courses.stream().map(Course::instructor).collect(Collectors.toList());
+		}
+
+	}
+
+}

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingPrincipalMethodArgumentResolverTests.java

@@ -0,0 +1,213 @@
+/*
+ * Copyright 2002-2021 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.data.method.annotation.support;
+
+import java.lang.reflect.Method;
+import java.security.Principal;
+import java.time.Duration;
+import java.util.function.Function;
+
+import graphql.ExecutionResult;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+import reactor.util.context.Context;
+
+import org.springframework.context.support.StaticApplicationContext;
+import org.springframework.core.MethodParameter;
+import org.springframework.graphql.GraphQlResponse;
+import org.springframework.graphql.GraphQlSetup;
+import org.springframework.graphql.RequestInput;
+import org.springframework.graphql.data.method.annotation.QueryMapping;
+import org.springframework.graphql.data.method.annotation.SubscriptionMapping;
+import org.springframework.graphql.execution.ExecutionGraphQlService;
+import org.springframework.graphql.execution.ReactorContextManager;
+import org.springframework.graphql.security.SecurityContextThreadLocalAccessor;
+import org.springframework.lang.Nullable;
+import org.springframework.security.authentication.TestingAuthenticationToken;
+import org.springframework.security.core.Authentication;
+import org.springframework.security.core.context.ReactiveSecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextHolder;
+import org.springframework.security.core.context.SecurityContextImpl;
+import org.springframework.stereotype.Controller;
+import org.springframework.util.ClassUtils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@code @SchemaMapping} methods with a {@link Principal} argument.
+ *
+ * @author Rossen Stoyanchev
+ */
+public class SchemaMappingPrincipalMethodArgumentResolverTests {
+
+	private final PrincipalMethodArgumentResolver resolver = new PrincipalMethodArgumentResolver();
+
+	private final Authentication authentication = new TestingAuthenticationToken(new Object(), new Object());
+
+	private final Function<Context, Context> reactiveContextWriter = context ->
+			ReactiveSecurityContextHolder.withAuthentication(this.authentication);
+
+	private final Function<Context, Context> threadLocalContextWriter = context ->
+			ReactorContextManager.extractThreadLocalValues(new SecurityContextThreadLocalAccessor(), context);
+
+	private final GreetingController greetingController = new GreetingController();
+
+
+
+	@Test
+	void supportsParameter() {
+		Method method = ClassUtils.getMethod(SchemaMappingPrincipalMethodArgumentResolverTests.class, "handle", (Class<?>[]) null);
+		assertThat(this.resolver.supportsParameter(new MethodParameter(method, 0))).isTrue();
+		assertThat(this.resolver.supportsParameter(new MethodParameter(method, 1))).isTrue();
+		assertThat(this.resolver.supportsParameter(new MethodParameter(method, 2))).isFalse();
+	}
+
+
+	@Nested
+	class Query {
+
+		@ParameterizedTest
+		@ValueSource(strings = {"greetingString", "greetingMono"})
+		void resolveFromReactiveContext(String field) {
+			testQuery(field, reactiveContextWriter);
+		}
+
+		@ParameterizedTest
+		@ValueSource(strings = {"greetingString", "greetingMono"})
+		void resolveFromThreadLocalContext(String field) {
+			SecurityContextHolder.setContext(new SecurityContextImpl(authentication));
+			try {
+				testQuery(field, threadLocalContextWriter);
+			}
+			finally {
+				SecurityContextHolder.clearContext();
+			}
+		}
+
+		private void testQuery(String field, Function<Context, Context> contextWriter) {
+			Mono<ExecutionResult> resultMono = executeAsync(
+					"type Query { " + field + ": String }", "{ " + field + " }", contextWriter);
+
+			String greeting = GraphQlResponse.from(resultMono).toEntity(field, String.class);
+			assertThat(greeting).isEqualTo("Hello");
+			assertThat(greetingController.principal()).isSameAs(authentication);
+		}
+
+	}
+
+
+	@Nested
+	class Subscription {
+
+		@Test
+		void resolveFromReactiveContext() {
+			testSubscription(reactiveContextWriter);
+		}
+
+		@Test
+		void resolveFromThreadLocalContext() {
+			SecurityContextHolder.setContext(new SecurityContextImpl(authentication));
+			try {
+				testSubscription(threadLocalContextWriter);
+			}
+			finally {
+				SecurityContextHolder.clearContext();
+			}
+		}
+
+		private void testSubscription(Function<Context, Context> contextModifier) {
+			String field = "greetingSubscription";
+
+			Mono<ExecutionResult> resultMono = executeAsync(
+					"type Query { greeting: String } type Subscription { " + field + ": String }",
+					"subscription Greeting { " + field + " }",
+					contextModifier);
+
+			Flux<String> greetingFlux = GraphQlResponse.forSubscription(resultMono)
+					.map(response -> response.toEntity(field, String.class));
+
+			StepVerifier.create(greetingFlux).expectNext("Hello", "Hi").verifyComplete();
+			assertThat(greetingController.principal()).isSameAs(authentication);
+		}
+
+	}
+
+	private Mono<ExecutionResult> executeAsync(
+			String schema, String op, Function<Context, Context> contextWriter) {
+
+		StaticApplicationContext context = new StaticApplicationContext();
+		context.getBeanFactory().registerSingleton("greetingController", greetingController);
+		context.refresh();
+
+		AnnotatedControllerConfigurer configurer = new AnnotatedControllerConfigurer();
+		configurer.setApplicationContext(context);
+		configurer.afterPropertiesSet();
+
+		ExecutionGraphQlService graphQlService =
+				GraphQlSetup.schemaContent(schema).runtimeWiring(configurer).toGraphQlService();
+
+		return Mono.delay(Duration.ofMillis(10))
+				.flatMap(aLong -> graphQlService.execute(new RequestInput(op, null, null, null)))
+				.contextWrite(contextWriter);
+	}
+
+
+	@SuppressWarnings("unused")
+	public void handle(
+			Principal principal,
+			Authentication authentication,
+			String s) {
+	}
+
+
+	@Controller
+	@SuppressWarnings("unused")
+	private static class GreetingController {
+
+		@Nullable
+		private Principal principal;
+
+		@Nullable
+		public Principal principal() {
+			return this.principal;
+		}
+
+		@QueryMapping
+		String greetingString(Principal principal) {
+			this.principal = principal;
+			return "Hello";
+		}
+
+		@QueryMapping
+		Mono<String> greetingMono(Principal principal) {
+			this.principal = principal;
+			return Mono.just("Hello");
+		}
+
+		@SubscriptionMapping
+		Flux<String> greetingSubscription(Principal principal) {
+			this.principal = principal;
+			return Flux.just("Hello", "Hi");
+		}
+
+	}
+
+}

spring-graphql/src/test/resources/log4j2-test.xml

@@ -8,6 +8,7 @@
 	<Loggers>
 		<Logger name="org.springframework" level="debug" />
 		<Logger name="org.springframework.graphql" level="trace" />
+		<Logger name="graphql" level="info" />
 		<Root level="error">
 			<AppenderRef ref="Console" />
 		</Root>