Add TypeVisitorHelper

See gh-722

Author: rstoyanchev

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

@@ -17,7 +17,7 @@
 package org.springframework.graphql.execution;
 
 import java.util.ArrayList;
-import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
@@ -132,17 +132,19 @@ private GraphQLSchema applyTypeVisitorsToTransformSchema(GraphQLSchema schema) {
 	}
 
 	private GraphQLSchema applyTypeVisitors(GraphQLSchema schema) {
-		GraphQLTypeVisitor visitor = ContextDataFetcherDecorator.createVisitor(schema, this.subscriptionExceptionResolvers);
-		List<GraphQLTypeVisitor> visitors = new ArrayList<>(this.typeVisitors);
-		visitors.add(visitor);
 
-		GraphQLCodeRegistry.Builder codeRegistry = GraphQLCodeRegistry.newCodeRegistry(schema.getCodeRegistry());
-		Map<Class<?>, Object> vars = Collections.singletonMap(GraphQLCodeRegistry.Builder.class, codeRegistry);
+		GraphQLCodeRegistry.Builder outputCodeRegistry =
+				GraphQLCodeRegistry.newCodeRegistry(schema.getCodeRegistry());
 
-		SchemaTraverser traverser = new SchemaTraverser();
-		traverser.depthFirstFullSchema(visitors, schema, vars);
+		Map<Class<?>, Object> vars = new HashMap<>(2);
+		vars.put(GraphQLCodeRegistry.Builder.class, outputCodeRegistry);
+		vars.put(TypeVisitorHelper.class, TypeVisitorHelper.create(schema));
 
-		return schema.transformWithoutTypes(builder -> builder.codeRegistry(codeRegistry));
+		List<GraphQLTypeVisitor> visitorsToUse = new ArrayList<>(this.typeVisitors);
+		visitorsToUse.add(ContextDataFetcherDecorator.createVisitor(this.subscriptionExceptionResolvers));
+
+		new SchemaTraverser().depthFirstFullSchema(visitorsToUse, schema, vars);
+		return schema.transformWithoutTypes(builder -> builder.codeRegistry(outputCodeRegistry));
 	}
 
 

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

@@ -25,8 +25,6 @@
 import graphql.schema.GraphQLCodeRegistry;
 import graphql.schema.GraphQLFieldDefinition;
 import graphql.schema.GraphQLFieldsContainer;
-import graphql.schema.GraphQLObjectType;
-import graphql.schema.GraphQLSchema;
 import graphql.schema.GraphQLSchemaElement;
 import graphql.schema.GraphQLTypeVisitor;
 import graphql.schema.GraphQLTypeVisitorStub;
@@ -114,25 +112,24 @@ public Object get(DataFetchingEnvironment environment) throws Exception {
 	 * Static factory method to create {@link GraphQLTypeVisitor} that wraps
 	 * data fetchers with the {@link ContextDataFetcherDecorator}.
 	 */
-	static GraphQLTypeVisitor createVisitor(
-			GraphQLSchema schema, List<SubscriptionExceptionResolver> resolvers) {
-
-		GraphQLObjectType subscriptionType = schema.getSubscriptionType();
-		String subscriptionTypeName = (subscriptionType != null ? subscriptionType.getName() : null);
+	static GraphQLTypeVisitor createVisitor(List<SubscriptionExceptionResolver> resolvers) {
 
 		SubscriptionExceptionResolver exceptionResolver = new CompositeSubscriptionExceptionResolver(resolvers);
 
 		return new GraphQLTypeVisitorStub() {
+
 			@Override
 			public TraversalControl visitGraphQLFieldDefinition(
 					GraphQLFieldDefinition fieldDefinition, TraverserContext<GraphQLSchemaElement> context) {
 
+				TypeVisitorHelper visitorHelper = context.getVarFromParents(TypeVisitorHelper.class);
 				GraphQLCodeRegistry.Builder codeRegistry = context.getVarFromParents(GraphQLCodeRegistry.Builder.class);
+
 				GraphQLFieldsContainer parent = (GraphQLFieldsContainer) context.getParentNode();
 				DataFetcher<?> dataFetcher = codeRegistry.getDataFetcher(parent, fieldDefinition);
 
 				if (applyDecorator(dataFetcher)) {
-					boolean handlesSubscription = parent.getName().equals(subscriptionTypeName);
+					boolean handlesSubscription = visitorHelper.isSubscriptionType(parent);
 					dataFetcher = new ContextDataFetcherDecorator(dataFetcher, handlesSubscription, exceptionResolver);
 					codeRegistry.dataFetcher(parent, fieldDefinition, dataFetcher);
 				}

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2002-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.execution;
+
+import graphql.schema.GraphQLNamedType;
+import graphql.schema.GraphQLObjectType;
+import graphql.schema.GraphQLSchema;
+
+import org.springframework.lang.Nullable;
+
+/**
+ * Default implementation of {@link TypeVisitorHelper} that performs checks
+ * against {@link GraphQLSchema}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.2.1
+ */
+final class DefaultTypeVisitorHelper implements TypeVisitorHelper {
+
+	@Nullable
+	private final String subscriptionTypeName;
+
+
+	/**
+	 * Package private constructor
+	 */
+	DefaultTypeVisitorHelper(GraphQLSchema schema) {
+		GraphQLObjectType subscriptionType = schema.getSubscriptionType();
+		this.subscriptionTypeName = (subscriptionType != null ? subscriptionType.getName() : null);
+	}
+
+
+	/**
+	 * Whether the given type is the subscription type.
+	 */
+	@Override
+	public boolean isSubscriptionType(GraphQLNamedType type) {
+		return (type.getName().equals(this.subscriptionTypeName));
+	}
+
+}

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2002-2023 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.graphql.execution;
+
+import java.util.List;
+
+import graphql.schema.GraphQLNamedType;
+import graphql.schema.GraphQLSchema;
+
+/**
+ * Helper for {@link graphql.schema.GraphQLTypeVisitor}s registered via
+ * {@link GraphQlSource.Builder#typeVisitors(List)} that is exposed as a
+ * variable in {@link graphql.util.TraverserContext}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 1.2.1
+ */
+public interface TypeVisitorHelper {
+
+	/**
+	 * Whether the given type is the subscription type.
+	 */
+	boolean isSubscriptionType(GraphQLNamedType type);
+
+
+	/**
+	 * Create an instance with the given {@link GraphQLSchema}.
+	 */
+	static TypeVisitorHelper create(GraphQLSchema schema) {
+		return new DefaultTypeVisitorHelper(schema);
+	}
+
+}