From 803896b2f453cc83437c65d445e1874416eede54 Mon Sep 17 00:00:00 2001
From: bbaker <bbaker@atlassian.com>
Date: Mon, 22 Dec 2025 14:26:13 +1100
Subject: [PATCH] An example of how we might fix the schema transformer

---
 .../graphql/schema/SchemaTransformer.java     | 35 +++++++++++--------
 .../schema/SchemaTransformerTest.groovy       | 14 ++++----
 2 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/src/main/java/graphql/schema/SchemaTransformer.java b/src/main/java/graphql/schema/SchemaTransformer.java
index 41d669f4e..801c66fa3 100644
--- a/src/main/java/graphql/schema/SchemaTransformer.java
+++ b/src/main/java/graphql/schema/SchemaTransformer.java
@@ -13,16 +13,7 @@
 import graphql.util.TraverserContext;
 import graphql.util.TraverserVisitor;
 
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Iterator;
-import java.util.LinkedHashMap;
-import java.util.LinkedHashSet;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.function.Consumer;
 
 import static graphql.Assert.assertNotEmpty;
@@ -539,7 +530,7 @@ private static class DummyRoot implements GraphQLSchemaElement {
         GraphQLObjectType mutation;
         GraphQLObjectType subscription;
         GraphQLObjectType introspectionSchemaType;
-        Set<GraphQLType> additionalTypes;
+        Set<GraphQLType> allTypesExceptOperationTypes;
         Set<GraphQLDirective> directives;
         Set<GraphQLDirective> schemaDirectives;
         Set<GraphQLAppliedDirective> schemaAppliedDirectives;
@@ -550,13 +541,27 @@ private static class DummyRoot implements GraphQLSchemaElement {
             query = schema.getQueryType();
             mutation = schema.isSupportingMutations() ? schema.getMutationType() : null;
             subscription = schema.isSupportingSubscriptions() ? schema.getSubscriptionType() : null;
-            additionalTypes = schema.getAdditionalTypes();
+            allTypesExceptOperationTypes = allTypesExceptOpOnes(schema);
             schemaDirectives = new LinkedHashSet<>(schema.getSchemaDirectives());
             schemaAppliedDirectives = new LinkedHashSet<>(schema.getSchemaAppliedDirectives());
             directives = new LinkedHashSet<>(schema.getDirectives());
             introspectionSchemaType = schema.getIntrospectionSchemaType();
         }
 
+        private Set<GraphQLType> allTypesExceptOpOnes(GraphQLSchema schema) {
+            Set<GraphQLType> types = new HashSet<>(schema.getTypeMap().values());
+            removeIfNotNull(types, schema.getQueryType());
+            removeIfNotNull(types, schema.getMutationType());
+            removeIfNotNull(types, schema.getSubscriptionType());
+            return types;
+        }
+
+        private void removeIfNotNull(Set<GraphQLType> types, GraphQLType type) {
+            if (type != null) {
+                types.remove(type);
+            }
+        }
+
         DummyRoot(GraphQLSchemaElement schemaElement) {
             this.schemaElement = schemaElement;
         }
@@ -584,7 +589,7 @@ public SchemaElementChildrenContainer getChildrenWithTypeReferences() {
                 if (schema.isSupportingSubscriptions()) {
                     builder.child(SUBSCRIPTION, subscription);
                 }
-                builder.children(ADD_TYPES, additionalTypes);
+                builder.children(ADD_TYPES, allTypesExceptOperationTypes);
                 builder.children(DIRECTIVES, directives);
                 builder.children(SCHEMA_DIRECTIVES, schemaDirectives);
                 builder.children(SCHEMA_APPLIED_DIRECTIVES, schemaAppliedDirectives);
@@ -604,7 +609,7 @@ public GraphQLSchemaElement withNewChildren(SchemaElementChildrenContainer newCh
             mutation = newChildren.getChildOrNull(MUTATION);
             subscription = newChildren.getChildOrNull(SUBSCRIPTION);
             introspectionSchemaType = newChildren.getChildOrNull(INTROSPECTION);
-            additionalTypes = new LinkedHashSet<>(newChildren.getChildren(ADD_TYPES));
+            allTypesExceptOperationTypes = new LinkedHashSet<>(newChildren.getChildren(ADD_TYPES));
             directives = new LinkedHashSet<>(newChildren.getChildren(DIRECTIVES));
             schemaDirectives = new LinkedHashSet<>(newChildren.getChildren(SCHEMA_DIRECTIVES));
             schemaAppliedDirectives = new LinkedHashSet<>(newChildren.getChildren(SCHEMA_APPLIED_DIRECTIVES));
@@ -621,7 +626,7 @@ public GraphQLSchema rebuildSchema(GraphQLCodeRegistry.Builder codeRegistry) {
                     .query(this.query)
                     .mutation(this.mutation)
                     .subscription(this.subscription)
-                    .additionalTypes(this.additionalTypes)
+                    .additionalTypes(this.allTypesExceptOperationTypes)
                     .additionalDirectives(this.directives)
                     .introspectionSchemaType(this.introspectionSchemaType)
                     .withSchemaDirectives(this.schemaDirectives)
diff --git a/src/test/groovy/graphql/schema/SchemaTransformerTest.groovy b/src/test/groovy/graphql/schema/SchemaTransformerTest.groovy
index 9513ed39f..5b2ec6f95 100644
--- a/src/test/groovy/graphql/schema/SchemaTransformerTest.groovy
+++ b/src/test/groovy/graphql/schema/SchemaTransformerTest.groovy
@@ -1056,13 +1056,13 @@ type Query {
         """
 
         def schema = TestUtil.schema(sdl)
-        schema = schema.transform { builder ->
-            for (def type : schema.getTypeMap().values()) {
-                if (type != schema.getQueryType() && type != schema.getMutationType() && type != schema.getSubscriptionType()) {
-                    builder.additionalType(type)
-                }
-            }
-        }
+//        schema = schema.transform { builder ->
+//            for (def type : schema.getTypeMap().values()) {
+//                if (type != schema.getQueryType() && type != schema.getMutationType() && type != schema.getSubscriptionType()) {
+//                    builder.additionalType(type)
+//                }
+//            }
+//        }
 
 
         def visitor = new GraphQLTypeVisitorStub() {