From f0ce9424bd31b4a275425dbe8f5d8ec2905bd9d3 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Tue, 20 Jul 2021 16:49:32 +0100
Subject: [PATCH] Apply type visitors via SchemaTraverser

Transforming the schema can be expensive and should be separated from
use cases where type visitors don't need to change the schema.

For now we'll provide read-only traversal only. Separately we can
provide an option for schema transformation when it becomes necessary.
---
 .../DefaultGraphQlSourceBuilder.java          | 28 ++++++++++++++-----
 .../graphql/execution/GraphQlSource.java      |  6 ++--
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultGraphQlSourceBuilder.java b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultGraphQlSourceBuilder.java
index 14e749f6..63772b69 100644
--- a/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultGraphQlSourceBuilder.java
+++ b/spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultGraphQlSourceBuilder.java
@@ -20,15 +20,18 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Consumer;
 
 import graphql.GraphQL;
 import graphql.execution.instrumentation.ChainedInstrumentation;
 import graphql.execution.instrumentation.Instrumentation;
+import graphql.schema.GraphQLCodeRegistry;
 import graphql.schema.GraphQLSchema;
 import graphql.schema.GraphQLTypeVisitor;
-import graphql.schema.SchemaTransformer;
+import graphql.schema.SchemaTraverser;
 import graphql.schema.idl.RuntimeWiring;
 import graphql.schema.idl.SchemaGenerator;
 import graphql.schema.idl.SchemaParser;
@@ -104,18 +107,14 @@ class DefaultGraphQlSourceBuilder implements GraphQlSource.Builder {
 				.orElseThrow(() -> new IllegalArgumentException("'schemaResources' should not be empty"));
 
 		GraphQLSchema schema = new SchemaGenerator().makeExecutableSchema(registry, this.runtimeWiring);
-		for (GraphQLTypeVisitor visitor : this.typeVisitors) {
-			schema = SchemaTransformer.transformSchema(schema, visitor);
-		}
-
-		// This comes last, wraps other DataFetcher's
-		schema = SchemaTransformer.transformSchema(schema, ContextDataFetcherDecorator.TYPE_VISITOR);
+		schema = applyTypeVisitors(schema);
 
 		GraphQL.Builder builder = GraphQL.newGraphQL(schema);
 		builder.defaultDataFetcherExceptionHandler(new ExceptionResolversExceptionHandler(this.exceptionResolvers));
 		if (!this.instrumentations.isEmpty()) {
 			builder = builder.instrumentation(new ChainedInstrumentation(this.instrumentations));
 		}
+
 		this.graphQlConfigurers.accept(builder);
 		GraphQL graphQl = builder.build();
 
@@ -135,6 +134,21 @@ class DefaultGraphQlSourceBuilder implements GraphQlSource.Builder {
 		}
 	}
 
+	private GraphQLSchema applyTypeVisitors(GraphQLSchema schema) {
+		List<GraphQLTypeVisitor> visitors = new ArrayList<>(this.typeVisitors);
+		visitors.add(ContextDataFetcherDecorator.TYPE_VISITOR);
+
+		GraphQLCodeRegistry.Builder builder = GraphQLCodeRegistry.newCodeRegistry(schema.getCodeRegistry());
+		Map<Class<?>, Object> vars = Collections.singletonMap(GraphQLCodeRegistry.Builder.class, builder);
+
+		SchemaTraverser traverser = new SchemaTraverser();
+		traverser.depthFirstFullSchema(visitors, schema, vars);
+
+		return GraphQLSchema.newSchema(schema).codeRegistry(builder.build()).build();
+	}
+
+
+
 	/**
 	 * GraphQlSource that returns the built GraphQL instance and its schema.
 	 */
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 1f40e2a2..d4625532 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
@@ -96,8 +96,10 @@ public interface GraphQlSource {
 		Builder exceptionResolvers(List<DataFetcherExceptionResolver> resolvers);
 
 		/**
-		 * Add {@link GraphQLTypeVisitor}'s to transform the underlying
-		 * {@link graphql.schema.GraphQLSchema} with.
+		 * Add {@link GraphQLTypeVisitor}s to visit all element of the created
+		 * {@link graphql.schema.GraphQLSchema}.
+		 * <p><strong>Note:</strong> Visitors are applied via
+		 * {@link graphql.schema.SchemaTraverser} and cannot change the schema.
 		 * @param typeVisitors the type visitors
 		 * @return the current builder
 		 * @see graphql.schema.SchemaTransformer#transformSchema(GraphQLSchema,