Do not skip DataFetcherFactories

Closes gh-440

Author: rstoyanchev

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

@@ -109,14 +109,14 @@ static GraphQLTypeVisitor createVisitor(List<SubscriptionExceptionResolver> reso
 
 		return new GraphQLTypeVisitorStub() {
 			@Override
-			public TraversalControl visitGraphQLFieldDefinition(GraphQLFieldDefinition fieldDefinition,
-																TraverserContext<GraphQLSchemaElement> context) {
+			public TraversalControl visitGraphQLFieldDefinition(
+					GraphQLFieldDefinition fieldDefinition, TraverserContext<GraphQLSchemaElement> context) {
 
 				GraphQLCodeRegistry.Builder codeRegistry = context.getVarFromParents(GraphQLCodeRegistry.Builder.class);
 				GraphQLFieldsContainer parent = (GraphQLFieldsContainer) context.getParentNode();
 				DataFetcher<?> dataFetcher = codeRegistry.getDataFetcher(parent, fieldDefinition);
 
-				if (dataFetcher.getClass().getPackage().getName().startsWith("graphql.")) {
+				if (skipDataFetcher(dataFetcher)) {
 					return TraversalControl.CONTINUE;
 				}
 
@@ -125,6 +125,14 @@ public TraversalControl visitGraphQLFieldDefinition(GraphQLFieldDefinition field
 				codeRegistry.dataFetcher(parent, fieldDefinition, dataFetcher);
 				return TraversalControl.CONTINUE;
 			}
+
+			private boolean skipDataFetcher(DataFetcher<?> dataFetcher) {
+				Class<?> type = dataFetcher.getClass();
+				if (type.getPackage().getName().startsWith("graphql.")) {
+					return !type.getSimpleName().startsWith("DataFetcherFactories");
+				}
+				return false;
+			}
 		};
 	}
 

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

@@ -19,12 +19,19 @@
 import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.BiConsumer;
 
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import graphql.GraphQL;
 import graphql.GraphQLError;
 import graphql.GraphqlErrorBuilder;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetcherFactories;
+import graphql.schema.GraphQLFieldDefinition;
+import graphql.schema.idl.SchemaDirectiveWiring;
+import graphql.schema.idl.SchemaDirectiveWiringEnvironment;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -42,10 +49,18 @@
  * Tests for {@link ContextDataFetcherDecorator}.
  * @author Rossen Stoyanchev
  */
+@SuppressWarnings("ReactiveStreamsUnusedPublisher")
 public class ContextDataFetcherDecoratorTests {
 
-	private static final String SCHEMA_CONTENT =
-			"type Query { greeting: String, greetings: [String] } type Subscription { greetings: String }";
+	private static final String SCHEMA_CONTENT = "" +
+			"directive @UpperCase on FIELD_DEFINITION " +
+			"type Query { " +
+			"  greeting: String @UpperCase, " +
+			"  greetings: [String] " +
+			"} " +
+			"type Subscription { " +
+			"  greetings: String " +
+			"}";
 
 
 	@Test
@@ -112,7 +127,7 @@ void fluxDataFetcherSubscription() throws Exception {
 	}
 
 	@Test
-	void fluxDataFetcherSubscriptionThrowException() throws Exception {
+	void fluxDataFetcherSubscriptionThrowingException() throws Exception {
 
 		SubscriptionExceptionResolver resolver =
 				SubscriptionExceptionResolver.forSingleError(exception ->
@@ -177,4 +192,51 @@ void dataFetcherWithThreadLocalContext() {
 		}
 	}
 
+	@Test // gh-440
+	void dataFetcherDecoratedWithDataFetcherFactories() {
+
+		SchemaDirectiveWiring directiveWiring = new SchemaDirectiveWiring() {
+
+			@SuppressWarnings("unchecked")
+			@Override
+			public GraphQLFieldDefinition onField(SchemaDirectiveWiringEnvironment<GraphQLFieldDefinition> env) {
+				if (env.getDirective("UpperCase") != null) {
+					return env.setFieldDataFetcher(DataFetcherFactories.wrapDataFetcher(
+							env.getFieldDataFetcher(),
+							((dataFetchingEnv, value) -> {
+								if (value instanceof String) {
+									return ((String) value).toUpperCase();
+								}
+								else if (value instanceof Mono) {
+									return ((Mono<String>) value).map(String::toUpperCase);
+								}
+								else {
+									throw new IllegalArgumentException();
+								}
+							})));
+				}
+				else {
+					return env.getElement();
+				}
+			}
+		};
+
+		BiConsumer<SchemaDirectiveWiring, DataFetcher<?>> tester = (schemaDirectiveWiring, dataFetcher) -> {
+
+			GraphQL graphQl = GraphQlSetup.schemaContent(SCHEMA_CONTENT)
+					.queryFetcher("greeting", dataFetcher)
+					.runtimeWiring(builder -> builder.directiveWiring(directiveWiring))
+					.toGraphQl();
+
+			ExecutionInput input = ExecutionInput.newExecutionInput().query("{ greeting }").build();
+			Mono<ExecutionResult> resultMono = Mono.fromFuture(graphQl.executeAsync(input));
+
+			String greeting = ResponseHelper.forResult(resultMono).toEntity("greeting", String.class);
+			assertThat(greeting).isEqualTo("HELLO");
+		};
+
+		tester.accept(directiveWiring, env -> CompletableFuture.completedFuture("hello"));
+		tester.accept(directiveWiring, env -> Mono.just("hello"));
+	}
+
 }