From fc042497f0d4fecdf9a78032591773ed48312a60 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Fri, 1 Oct 2021 19:17:24 +0100
Subject: [PATCH] Polishing contribution

Closes gh-153
---
 .../DefaultGraphQlSourceBuilder.java          | 22 ++++++++++++-------
 .../graphql/execution/GraphQlSource.java      | 19 +++++++++-------
 2 files changed, 25 insertions(+), 16 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 720a7895..0ec3030b 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
@@ -39,6 +39,7 @@ import graphql.schema.idl.SchemaParser;
 import graphql.schema.idl.TypeDefinitionRegistry;
 
 import org.springframework.core.io.Resource;
+import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 
 /**
@@ -60,12 +61,12 @@ class DefaultGraphQlSourceBuilder implements GraphQlSource.Builder {
 
 	private final List<Instrumentation> instrumentations = new ArrayList<>();
 
+	@Nullable
+	private BiFunction<TypeDefinitionRegistry, RuntimeWiring, GraphQLSchema> schemaFactory;
+
 	private Consumer<GraphQL.Builder> graphQlConfigurers = (builder) -> {
 	};
 	
-	private BiFunction<TypeDefinitionRegistry, RuntimeWiring, GraphQLSchema> schemaFactory =
-			(typeRegistry, runtimeWiring) -> new SchemaGenerator().makeExecutableSchema(typeRegistry, runtimeWiring);
-
 
 	@Override
 	public GraphQlSource.Builder schemaResources(Resource... resources) {
@@ -98,14 +99,16 @@ class DefaultGraphQlSourceBuilder implements GraphQlSource.Builder {
 	}
 
 	@Override
-	public GraphQlSource.Builder configureGraphQl(Consumer<GraphQL.Builder> configurer) {
-		this.graphQlConfigurers = this.graphQlConfigurers.andThen(configurer);
+	public GraphQlSource.Builder schemaFactory(
+			BiFunction<TypeDefinitionRegistry, RuntimeWiring, GraphQLSchema> schemaFactory) {
+
+		this.schemaFactory = schemaFactory;
 		return this;
 	}
 
 	@Override
-	public GraphQlSource.Builder schemaFactory(BiFunction<TypeDefinitionRegistry, RuntimeWiring, GraphQLSchema> schemaFactory) {
-		this.schemaFactory = schemaFactory;
+	public GraphQlSource.Builder configureGraphQl(Consumer<GraphQL.Builder> configurer) {
+		this.graphQlConfigurers = this.graphQlConfigurers.andThen(configurer);
 		return this;
 	}
 
@@ -117,8 +120,11 @@ class DefaultGraphQlSourceBuilder implements GraphQlSource.Builder {
 
 		RuntimeWiring.Builder runtimeWiringBuilder = RuntimeWiring.newRuntimeWiring();
 		this.runtimeWiringConfigurers.forEach(configurer -> configurer.configure(runtimeWiringBuilder));
+		RuntimeWiring runtimeWiring = runtimeWiringBuilder.build();
 
-		GraphQLSchema schema = schemaFactory.apply(registry, runtimeWiringBuilder.build());
+		GraphQLSchema schema = (this.schemaFactory != null ?
+				this.schemaFactory.apply(registry, runtimeWiring) :
+				new SchemaGenerator().makeExecutableSchema(registry, runtimeWiring));
 		schema = applyTypeVisitors(schema);
 
 		GraphQL.Builder builder = GraphQL.newGraphQL(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 bb1111f9..e853f8d1 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
@@ -118,6 +118,17 @@ public interface GraphQlSource {
 		 */
 		Builder instrumentation(List<Instrumentation> instrumentations);
 
+		/**
+		 * Configure a function to create the {@link GraphQLSchema} instance from the
+		 * given {@link TypeDefinitionRegistry} and {@link RuntimeWiring}. This may
+		 * be useful for federation to create a combined schema.
+		 * <p>By default, the schema is created with
+		 * {@link graphql.schema.idl.SchemaGenerator#makeExecutableSchema}.
+		 * @param schemaFactory the function to create the schema
+		 * @return the current builder
+		 */
+		Builder schemaFactory(BiFunction<TypeDefinitionRegistry, RuntimeWiring, GraphQLSchema> schemaFactory);
+
 		/**
 		 * Configure consumers to be given access to the {@link GraphQL.Builder} used to
 		 * build {@link GraphQL}.
@@ -126,14 +137,6 @@ public interface GraphQlSource {
 		 */
 		Builder configureGraphQl(Consumer<GraphQL.Builder> configurer);
 
-		/**
-		 * Provide a custom factory for creating an executable {@link GraphQLSchema} given
-		 * a {@link TypeDefinitionRegistry} and a {@link RuntimeWiring}
-		 * @param schemaFactory the schema factory
-		 * @return the current builder
-		 */
-		Builder schemaFactory(BiFunction<TypeDefinitionRegistry, RuntimeWiring, GraphQLSchema> schemaFactory);
-
 		/**
 		 * Build the {@link GraphQlSource}.
 		 * @return the built GraphQlSource