From bd3488899635bffcf53569ddd05f61156ef5be46 Mon Sep 17 00:00:00 2001
From: Robert Mosolgo <rdmosolgo@gmail.com>
Date: Wed, 4 Sep 2024 13:59:38 -0400
Subject: [PATCH] fix printing union possible types when the union is an extra
 type

---
 .../document_from_schema_definition.rb        | 33 ++++++++++++++++---
 spec/graphql/schema/printer_spec.rb           | 26 +++++++++++++++
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/lib/graphql/language/document_from_schema_definition.rb b/lib/graphql/language/document_from_schema_definition.rb
index 6dd6c43773..ff00c71605 100644
--- a/lib/graphql/language/document_from_schema_definition.rb
+++ b/lib/graphql/language/document_from_schema_definition.rb
@@ -18,6 +18,7 @@ def initialize(
         include_built_in_directives: false, include_built_in_scalars: false, always_include_schema: false
       )
         @schema = schema
+        @context = context
         @always_include_schema = always_include_schema
         @include_introspection_types = include_introspection_types
         @include_built_in_scalars = include_built_in_scalars
@@ -268,7 +269,33 @@ def build_definition_nodes
         end
         definitions = build_directive_nodes(dirs_to_build)
         all_types = @types.all_types
-        type_nodes = build_type_definition_nodes(all_types + schema.extra_types)
+        type_nodes = build_type_definition_nodes(all_types)
+
+        if (ex_t = schema.extra_types).any?
+          dummy_query = Class.new(GraphQL::Schema::Object) do
+            graphql_name "DummyQuery"
+            (all_types + ex_t).each_with_index do |type, idx|
+              if !type.kind.input_object? && !type.introspection?
+                field "f#{idx}", type
+              end
+            end
+          end
+
+          extra_types_schema = Class.new(GraphQL::Schema) do
+            query(dummy_query)
+          end
+
+          extra_types_types = GraphQL::Query.new(extra_types_schema, "{ __typename }", context: @context).types  # rubocop:disable Development/ContextIsPassedCop
+          # Temporarily replace `@types` with something from this example schema.
+          # It'd be much nicer to pass this in, but that would be a big refactor :S
+          prev_types = @types
+          @types = extra_types_types
+          type_nodes += build_type_definition_nodes(ex_t)
+          @types = prev_types
+        end
+
+        type_nodes.sort_by!(&:name)
+
         if @include_one_of
           # This may have been set to true when iterating over all types
           definitions.concat(build_directive_nodes([GraphQL::Schema::Directive::OneOf]))
@@ -291,9 +318,7 @@ def build_type_definition_nodes(types)
           types = types.reject { |type| type.kind.scalar? && type.default_scalar? }
         end
 
-        types
-          .map { |type| build_type_definition_node(type) }
-          .sort_by(&:name)
+        types.map { |type| build_type_definition_node(type) }
       end
 
       def build_field_nodes(fields)
diff --git a/spec/graphql/schema/printer_spec.rb b/spec/graphql/schema/printer_spec.rb
index 7906daf11f..938c4c8e43 100644
--- a/spec/graphql/schema/printer_spec.rb
+++ b/spec/graphql/schema/printer_spec.rb
@@ -967,4 +967,30 @@ class OddlyNamedQuery < GraphQL::Schema::Object
     schema = GraphQL::Schema.from_definition(input)
     assert_equal input, GraphQL::Schema::Printer.print_schema(schema)
   end
+
+  describe "when Union is used in extra_types" do
+    it "can be included" do
+      obj_1 = Class.new(GraphQL::Schema::Object) { graphql_name("Obj1"); field(:f1, String)}
+      obj_2 = Class.new(GraphQL::Schema::Object) { graphql_name("Obj2"); field(:f2, obj_1) }
+      union_type = Class.new(GraphQL::Schema::Union) do
+        graphql_name "Union1"
+        possible_types(obj_1, obj_2)
+      end
+
+      assert_equal "union Union1 = Obj1 | Obj2\n", Class.new(GraphQL::Schema) { extra_types(union_type) }.to_definition
+
+      expected_defn = <<~GRAPHQL
+        type Obj1 {
+          f1: String
+        }
+
+        type Obj2 {
+          f2: Obj1
+        }
+
+        union Union1 = Obj1 | Obj2
+      GRAPHQL
+      assert_equal expected_defn, Class.new(GraphQL::Schema) { extra_types(union_type, obj_1, obj_2) }.to_definition
+    end
+  end
 end