Support GraphQLContext as a method argument

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

@@ -516,6 +516,10 @@ Annotated handler methods can choose from one of the following method arguments:
 | For direct access to the underlying `DataFetchingEnvironment`.
   See <<controllers-environment>>.
 
+| `GraphQLContext`
+| For access to the context from the `DataFetchingEnvironment`.
+See <<controllers-graphql-context>>.
+
 |===
 
 
@@ -590,8 +594,15 @@ explicitly specify the type name in the `@SchemaMapping` annotation.
 [[controllers-environment]]
 ==== `DataFetchingEnvironment`
 
-To access the `DataFetchingEnvironment` directly, simply declare a method parameter of
-the same type.
+To access the `DataFetchingEnvironment` directly, declare a method parameter of the same
+type.
+
+
+[[controllers-graphql-context]]
+==== `GraphQLContext`
+
+To access the `GraphQLContext` from the `DataFetchingEnvironment`, declare a method
+parameter of the same type.
 
 
 

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

@@ -15,13 +15,14 @@
  */
 package org.springframework.graphql.data.method.annotation.support;
 
+import graphql.GraphQLContext;
 import graphql.schema.DataFetchingEnvironment;
 
 import org.springframework.core.MethodParameter;
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
 
 /**
- * Resolver for a {@link DataFetchingEnvironment} argument.
+ * Resolver for {@link DataFetchingEnvironment} or {@link GraphQLContext} arguments.
  *
  * @author Rossen Stoyanchev
  * @since 1.0.0
@@ -30,12 +31,19 @@ public class DataFetchingEnvironmentMethodArgumentResolver implements HandlerMet
 
 	@Override
 	public boolean supportsParameter(MethodParameter parameter) {
-		return parameter.getParameterType().equals(DataFetchingEnvironment.class);
+		Class<?> type = parameter.getParameterType();
+		return (type.equals(DataFetchingEnvironment.class) || type.equals(GraphQLContext.class));
 	}
 
 	@Override
 	public Object resolveArgument(MethodParameter parameter, DataFetchingEnvironment environment) {
-		return environment;
+		Class<?> type = parameter.getParameterType();
+		if (type.equals(GraphQLContext.class)) {
+			return environment.getGraphQlContext();
+		}
+		else {
+			return environment;
+		}
 	}
 
 }

spring-graphql/src/test/java/org/springframework/graphql/data/method/AnnotatedDataFetcherInvocationTests.java

@@ -18,8 +18,10 @@ package org.springframework.graphql.data.method;
 import java.util.List;
 import java.util.Map;
 
+import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import graphql.GraphQL;
+import graphql.GraphQLContext;
 import graphql.schema.DataFetchingEnvironment;
 import org.junit.jupiter.api.Test;
 import org.reactivestreams.Publisher;
@@ -108,7 +110,8 @@ public class AnnotatedDataFetcherInvocationTests {
 				"  }" +
 				"}";
 
-		ExecutionResult result = initGraphQl(BookController.class).execute(query);
+		ExecutionInput input = ExecutionInput.newExecutionInput().query(query).build();
+		ExecutionResult result = initGraphQl(BookController.class).execute(input);
 
 		assertThat(result.getErrors()).isEmpty();
 		Map<String, Object> data = result.getData();
@@ -118,6 +121,8 @@ public class AnnotatedDataFetcherInvocationTests {
 		assertThat(author.get("id")).isEqualTo("1");
 		assertThat(author.get("firstName")).isEqualTo("George");
 		assertThat(author.get("lastName")).isEqualTo("Orwell");
+
+		assertThat(input.getGraphQLContext().<String>get("key")).isEqualTo("value");
 	}
 
 	@Test
@@ -218,7 +223,8 @@ public class AnnotatedDataFetcherInvocationTests {
 		}
 
 		@QueryMapping
-		public Author authorById(DataFetchingEnvironment environment) {
+		public Author authorById(DataFetchingEnvironment environment, GraphQLContext context) {
+			context.put("key", "value");
 			String id = environment.getArgument("id");
 			return BookSource.getAuthor(Long.parseLong(id));
 		}