Support schema transforming GraphQLTypeVisitor

Closes gh-536

Author: rstoyanchev

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

@@ -284,29 +284,23 @@ The main implementation, `DefaultExecutionGraphQlService`, is configured with a
 [[execution-graphqlsource]]
 === `GraphQLSource`
 
-`GraphQlSource` is a core Spring abstraction for access to the
-`graphql.GraphQL` instance to use for request execution. It provides a builder API to
-initialize GraphQL Java and build a `GraphQlSource`.
-
-The default `GraphQlSource` builder, accessible via
-`GraphQlSource.schemaResourceBuilder()`, enables support for
-<<execution-reactive-datafetcher>>, <<execution-context>>, and <<execution-exceptions>>.
-
-The Spring Boot {spring-boot-ref-docs}/web.html#web.graphql[starter] initializes a
-`GraphQlSource` instance through the default `GraphQlSource.Builder` and also enables
-the following:
-
-- Load <<execution-graphqlsource-schema-resources, schema files>> from a configurable location.
-- Expose {spring-boot-ref-docs}/application-properties.html#appendix.application-properties.web[properties]
+`GraphQlSource` is a contract to expose the `graphql.GraphQL` instance to use that also
+includes a builder API to build that instance. The default builder is available via
+`GraphQlSource.schemaResourceBuilder()`. The
+{spring-boot-ref-docs}/web.html#web.graphql[Spring Boot starter] creates an instance of
+this builder and further initializes it as follows:
+
+- Loads <<execution-graphqlsource-schema-resources, schema files>> from a configurable location.
+- Exposes {spring-boot-ref-docs}/application-properties.html#appendix.application-properties.web[properties]
   that apply to `GraphQlSource.Builder`.
-- Detect <<execution-graphqlsource-runtimewiring-configurer>> beans.
-- Detect https://www.graphql-java.com/documentation/instrumentation[Instrumentation] beans for
+- Detects <<execution-graphqlsource-runtimewiring-configurer>> beans.
+- Detects https://www.graphql-java.com/documentation/instrumentation[Instrumentation] beans for
   {spring-boot-ref-docs}/actuator.html#actuator.metrics.supported.spring-graphql[GraphQL metrics].
-- Detect `DataFetcherExceptionResolver` beans for <<execution-exceptions, exception resolution>>.
-- Detect `SubscriptionExceptionResolver` beans for <<execution-exceptions-subsctiption, subscription exception resolution>>.
+- Detects `DataFetcherExceptionResolver` beans for <<execution-exceptions, exception resolution>>.
+- Detects `SubscriptionExceptionResolver` beans for <<execution-exceptions-subsctiption, subscription exception resolution>>.
 
-For further customizations, you can declare your own `GraphQlSourceBuilderCustomizer` beans;
-for example, for configuring your own `ExecutionIdProvider`:
+For further customizations, you can declare a `GraphQlSourceBuilderCustomizer` bean. For example, to
+configure your own `ExecutionIdProvider`:
 
 [source,java,indent=0,subs="verbatim,quotes"]
 ----
@@ -345,9 +339,9 @@ locations, e.g. across multiple modules.
 [[execution-graphqlsource-schema-creation]]
 ==== Schema Creation
 
-By default, `GraphQlSource.Builder` uses the GraphQL Java `GraphQLSchemaGenerator` to
-create the `graphql.schema.GraphQLSchema`. This works for most applications, but if
-necessary, you can hook into the schema creation through the builder:
+By default, `GraphQlSource.Builder` uses the GraphQL Java `SchemaGenerator` to create the
+`graphql.schema.GraphQLSchema`. This works for typical use, but if you need to use a
+different generator, e.g. for federation, you can register a `schemaFactory` callback:
 
 [source,java,indent=0,subs="verbatim,quotes"]
 ----
@@ -360,10 +354,29 @@ builder.schemaResources(..)
 		})
 ----
 
-The primary reason for this is to create the schema through a federation library.
-
 The <<execution-graphqlsource, GraphQlSource section>> explains how to configure that with Spring Boot.
 
+
+[[execution-graphqlsource-schema-traversal]]
+==== Schema Traversal
+
+You can register a `graphql.schema.GraphQLTypeVisitor` via
+`builder.schemaResources(..).typeVisitors(..)` if you want to traverse the schema after
+it is created, and possibly apply changes to the `GraphQLCodeRegistry`. Keep in mind,
+however, that such a visitor cannot change the schema. See
+<<execution-graphqlsource-schema-transformation>>, if you need to make changes to the schema.
+
+
+[[execution-graphqlsource-schema-transformation]]
+==== Schema Transformation
+
+You can register a `graphql.schema.GraphQLTypeVisitor` via
+`builder.schemaResources(..).typeVisitorsToTransformSchema(..)` if you want to traverse
+and transform the schema after it is created, and make changes to the schema. Keep in mind
+that this is more expensive than <<execution-graphqlsource-schema-traversal>> so generally
+prefer traversal to transformation unless you need to make schema changes.
+
+
 [[execution-graphqlsource-runtimewiring-configurer]]
 ==== `RuntimeWiringConfigurer`
 

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

@@ -28,6 +28,7 @@
 import graphql.schema.GraphQLCodeRegistry;
 import graphql.schema.GraphQLSchema;
 import graphql.schema.GraphQLTypeVisitor;
+import graphql.schema.SchemaTransformer;
 import graphql.schema.SchemaTraverser;
 
 
@@ -47,6 +48,8 @@ abstract class AbstractGraphQlSourceBuilder<B extends GraphQlSource.Builder<B>>
 
 	private final List<GraphQLTypeVisitor> typeVisitors = new ArrayList<>();
 
+	private final List<GraphQLTypeVisitor> typeVisitorsToTransformSchema = new ArrayList<>();
+
 	private final List<Instrumentation> instrumentations = new ArrayList<>();
 
 	private Consumer<GraphQL.Builder> graphQlConfigurers = (builder) -> {
@@ -71,6 +74,12 @@ public B typeVisitors(List<GraphQLTypeVisitor> typeVisitors) {
 		return self();
 	}
 
+	@Override
+	public B typeVisitorsToTransformSchema(List<GraphQLTypeVisitor> typeVisitorsToTransformSchema) {
+		this.typeVisitorsToTransformSchema.addAll(typeVisitorsToTransformSchema);
+		return self();
+	}
+
 	@Override
 	public B instrumentation(List<Instrumentation> instrumentations) {
 		this.instrumentations.addAll(instrumentations);
@@ -92,6 +101,7 @@ private  <T extends B> T self() {
 	public GraphQlSource build() {
 		GraphQLSchema schema = initGraphQlSchema();
 
+		schema = applyTypeVisitorsToTransformSchema(schema);
 		schema = applyTypeVisitors(schema);
 
 		GraphQL.Builder builder = GraphQL.newGraphQL(schema);
@@ -112,6 +122,14 @@ public GraphQlSource build() {
 	 */
 	protected abstract GraphQLSchema initGraphQlSchema();
 
+	private GraphQLSchema applyTypeVisitorsToTransformSchema(GraphQLSchema schema) {
+		SchemaTransformer transformer = new SchemaTransformer();
+		for (GraphQLTypeVisitor visitor : this.typeVisitorsToTransformSchema) {
+			schema = transformer.transform(schema, visitor);
+		}
+		return schema;
+	}
+
 	private GraphQLSchema applyTypeVisitors(GraphQLSchema schema) {
 		GraphQLTypeVisitor visitor = ContextDataFetcherDecorator.createVisitor(this.subscriptionExceptionResolvers);
 		List<GraphQLTypeVisitor> visitors = new ArrayList<>(this.typeVisitors);

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

@@ -105,7 +105,8 @@ interface Builder<B extends Builder<B>> {
 
 		/**
 		 * Add {@link GraphQLTypeVisitor}s to visit all element of the created
-		 * {@link graphql.schema.GraphQLSchema}.
+		 * {@link graphql.schema.GraphQLSchema} and make changes to the
+		 * {@link graphql.schema.GraphQLCodeRegistry}.
 		 * <p><strong>Note:</strong> Visitors are applied via
 		 * {@link graphql.schema.SchemaTraverser} and cannot change the schema.
 		 * @param typeVisitors the type visitors
@@ -114,6 +115,21 @@ interface Builder<B extends Builder<B>> {
 		 */
 		B typeVisitors(List<GraphQLTypeVisitor> typeVisitors);
 
+		/**
+		 * Alternative to {@link #typeVisitors(List)} for visitors that also
+		 * need to make schema changes.
+		 * <p><strong>Note:</strong> Visitors are applied via
+		 * {@link graphql.schema.SchemaTransformer}, and therefore can change
+		 * the schema. However, this is more expensive than using
+		 * {@link graphql.schema.SchemaTraverser}, so generally prefer
+		 * {@link #typeVisitors(List)} if it's not necessary to change the schema.
+		 * @param typeVisitors the type visitors to register
+		 * @return the current builder
+		 * @see graphql.schema.SchemaTransformer#transformSchema(GraphQLSchema, GraphQLTypeVisitor)
+		 * @since 1.1
+		 */
+		B typeVisitorsToTransformSchema(List<GraphQLTypeVisitor> typeVisitors);
+
 		/**
 		 * Provide {@link Instrumentation} components to instrument the
 		 * execution of GraphQL queries.

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

@@ -16,14 +16,22 @@
 package org.springframework.graphql.execution;
 
 import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
 
+import graphql.Scalars;
 import graphql.schema.DataFetcher;
 import graphql.schema.FieldCoordinates;
 import graphql.schema.GraphQLFieldDefinition;
+import graphql.schema.GraphQLObjectType;
 import graphql.schema.GraphQLSchema;
+import graphql.schema.GraphQLSchemaElement;
+import graphql.schema.GraphQLTypeVisitor;
+import graphql.schema.GraphQLTypeVisitorStub;
 import graphql.schema.idl.FieldWiringEnvironment;
 import graphql.schema.idl.RuntimeWiring;
 import graphql.schema.idl.WiringFactory;
+import graphql.util.TraversalControl;
+import graphql.util.TraverserContext;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.graphql.BookSource;
@@ -45,6 +53,65 @@ void duplicateResourcesAreIgnored() {
 		GraphQlSetup.schemaResource(BookSource.schema, BookSource.schema).toGraphQlSource();
 	}
 
+	@Test
+	void typeVisitors() {
+
+		AtomicInteger counter = new AtomicInteger();
+
+		GraphQLTypeVisitor visitor = new GraphQLTypeVisitorStub() {
+
+			@Override
+			public TraversalControl visitGraphQLObjectType(
+					GraphQLObjectType node, TraverserContext<GraphQLSchemaElement> context) {
+
+				counter.incrementAndGet();
+				return TraversalControl.CONTINUE;
+			}
+		};
+
+		GraphQlSetup.schemaContent("type Query {  myQuery: String}").typeVisitor(visitor).toGraphQlSource();
+
+		assertThat(counter.get()).isPositive();
+	}
+
+	@Test
+	void typeVisitorToTransformSchema() {
+
+		String schemaContent = "" +
+				"type Query {" +
+				"  person: Person" +
+				"} " +
+				"type Person {" +
+				"  firstName: String" +
+				"}";
+
+		GraphQLTypeVisitor visitor = new GraphQLTypeVisitorStub() {
+
+			@Override
+			public TraversalControl visitGraphQLObjectType(
+					GraphQLObjectType node, TraverserContext<GraphQLSchemaElement> context) {
+
+				if (node.getName().equals("Person")) {
+					node = node.transform(builder -> builder.field(
+							GraphQLFieldDefinition.newFieldDefinition()
+									.name("lastName")
+									.type(Scalars.GraphQLString)
+									.build()));
+					changeNode(context, node);
+				}
+
+				return TraversalControl.CONTINUE;
+			}
+		};
+
+		GraphQLSchema schema = GraphQlSetup.schemaContent(schemaContent)
+				.typeVisitorToTransformSchema(visitor)
+				.toGraphQlSource()
+				.schema();
+
+		assertThat(schema.getObjectType("Person").getFieldDefinition("lastName")).isNotNull();
+	}
+
 	@Test
 	void wiringFactoryList() {
 

spring-graphql/src/testFixtures/java/org/springframework/graphql/GraphQlSetup.java

@@ -117,6 +117,11 @@ public GraphQlSetup typeVisitor(GraphQLTypeVisitor... visitors) {
 		return this;
 	}
 
+	public GraphQlSetup typeVisitorToTransformSchema(GraphQLTypeVisitor... visitors) {
+		this.graphQlSourceBuilder.typeVisitorsToTransformSchema(Arrays.asList(visitors));
+		return this;
+	}
+
 	public GraphQL toGraphQl() {
 		return this.graphQlSourceBuilder.build().graphQl();
 	}