Implement more visibility preload hooks

lib/graphql/schema.rb

@@ -512,10 +512,10 @@ def subscription(new_subscription_object = nil, &lazy_load_block)
             raise GraphQL::Error, "Second definition of `subscription(...)` (#{dup_defn.inspect}) is invalid, already configured with #{@subscription_object.inspect}"
           elsif use_visibility_profile?
             if block_given?
-              @subscription_object = lazy_load_block
+              @subscription_object = new_subscription_object
             else
-              @subscription_object = new_mutation_object
-              self.visibility.mutation_configured(@subscription_object)
+              @subscription_object = new_subscription_object
+              self.visibility.subscription_configured(@subscription_object)
             end
             add_subscription_extension_if_necessary
           else
@@ -717,6 +717,7 @@ def introspection(new_introspection_namespace = nil)
           @introspection = new_introspection_namespace
           # reset this cached value:
           @introspection_system = nil
+          introspection_system
         else
           @introspection || find_inherited_value(:introspection)
         end
@@ -726,6 +727,7 @@ def introspection_system
         if !@introspection_system
           @introspection_system = Schema::IntrospectionSystem.new(self)
           @introspection_system.resolve_late_bindings
+          self.visibility&.introspection_system_configured(@introspection_system)
         end
         @introspection_system
       end
@@ -969,6 +971,13 @@ def extra_types(*new_extra_types)
         end
       end
 
+      # Tell the schema about these types so that they can be registered as implementations of interfaces in the schema.
+      #
+      # This method must be used when an object type is connected to the schema as an interface implementor but
+      # not as a return type of a field. In that case, if the object type isn't registered here, GraphQL-Ruby won't be able to find it.
+      #
+      # @param new_orphan_types [Array<Class<GraphQL::Schema::Object>>] Object types to register as implementations of interfaces in the schema.
+      # @return [Array<Class<GraphQL::Schema::Object>>] All previously-registered orphan types for this schema
       def orphan_types(*new_orphan_types)
         if new_orphan_types.any?
           new_orphan_types = new_orphan_types.flatten
@@ -985,6 +994,7 @@ def orphan_types(*new_orphan_types)
           end
           add_type_and_traverse(new_orphan_types, root: false) unless use_visibility_profile?
           own_orphan_types.concat(new_orphan_types.flatten)
+          self.visibility&.orphan_types_configured(new_orphan_types)
         end
 
         inherited_ot = find_inherited_value(:orphan_types, nil)

lib/graphql/schema/visibility.rb

@@ -27,48 +27,67 @@ def initialize(schema, dynamic:, preload:, profiles:, migration_errors:)
         @dynamic = dynamic
         @migration_errors = migration_errors
         if preload
+          # Traverse the schema now (and in the *_configured hooks below)
+          # To make sure things are loaded during boot
           @preloaded_types = Set.new
-          if profiles.any?
-            profiles.each do |profile_name, example_ctx|
-              example_ctx[:visibility_profile] = profile_name
-              prof = profile_for(example_ctx, profile_name)
-              prof.all_types # force loading
-            end
-          else
-            types_to_visit = [
-              @schema.query,
-              @schema.mutation,
-              @schema.subscription,
-              *@schema.introspection_system.types.values,
-              *@schema.orphan_types,
-              *@schema.directives.map { |name, dir| dir.all_argument_definitions.map { |defn| defn.type.unwrap } }.flatten
-            ]
-            types_to_visit.compact!
-            ensure_all_loaded(types_to_visit)
-            # Somehow load files here
+          types_to_visit = [
+            @schema.query,
+            @schema.mutation,
+            @schema.subscription,
+            *@schema.introspection_system.types.values,
+            *@schema.introspection_system.entry_points.map { |ep| ep.type.unwrap },
+            *@schema.orphan_types,
+          ]
+          # Root types may have been nil:
+          types_to_visit.compact!
+          ensure_all_loaded(types_to_visit)
+
+          profiles.each do |profile_name, example_ctx|
+            example_ctx[:visibility_profile] = profile_name
+            prof = profile_for(example_ctx, profile_name)
+            prof.all_types # force loading
           end
         end
       end
 
+      # @api private
       def query_configured(query_type)
         if @preload
           ensure_all_loaded([query_type])
         end
       end
 
+      # @api private
       def mutation_configured(mutation_type)
         if @preload
           ensure_all_loaded([mutation_type])
         end
       end
 
+      # @api private
       def subscription_configured(subscription_type)
         if @preload
           ensure_all_loaded([subscription_type])
         end
       end
 
-      # TODO: on_orphan_types, on_directive, on_introspection_system
+      # @api private
+      def orphan_types_configured(orphan_types)
+        if @preload
+          ensure_all_loaded(orphan_types)
+        end
+      end
+
+      # @api private
+      def introspection_system_configured(introspection_system)
+        if @preload
+          introspection_types = [
+            *@schema.introspection_system.types.values,
+            *@schema.introspection_system.entry_points.map { |ep| ep.type.unwrap },
+          ]
+          ensure_all_loaded(introspection_types)
+        end
+      end
 
       # Make another Visibility for `schema` based on this one
       # @return [Visibility]

spec/graphql/schema/visibility_spec.rb

@@ -89,8 +89,8 @@ def exec_query(...)
 
     describe "when no profile is defined" do
       class NoProfileSchema < GraphQL::Schema
-        class ExampleExtension < GraphQL::Schema::FieldExtension
-        end
+        class ExampleExtension < GraphQL::Schema::FieldExtension; end
+        class OtherExampleExtension < GraphQL::Schema::FieldExtension; end
 
         class Query < GraphQL::Schema::Object
           field :str do
@@ -99,13 +99,54 @@ class Query < GraphQL::Schema::Object
           end
         end
 
+        class Mutation < GraphQL::Schema::Object
+          field :str do
+            type(String)
+            extension(ExampleExtension)
+          end
+        end
+
+        class Subscription < GraphQL::Schema::Object
+          field :str do
+            type(String)
+            extension(ExampleExtension)
+          end
+        end
+
+        class OrphanType < GraphQL::Schema::Object
+          field :str do
+            type(String)
+            extension(ExampleExtension)
+            extension(OtherExampleExtension)
+          end
+        end
+        # This one is added before `Visibility`
+        mutation(Mutation)
         use GraphQL::Schema::Visibility, preload: true
+        subscription(Subscription)
+        orphan_types(OrphanType)
+
+        module CustomIntrospection
+          class DynamicFields < GraphQL::Introspection::DynamicFields
+            field :__hello do
+              type(String)
+              extension(OtherExampleExtension)
+            end
+          end
+        end
       end
 
       it "still preloads" do
         assert_equal [], NoProfileSchema::Query.all_field_definitions.first.extensions.map(&:class)
         NoProfileSchema.query(NoProfileSchema::Query)
         assert_equal [NoProfileSchema::ExampleExtension], NoProfileSchema::Query.all_field_definitions.first.extensions.map(&:class)
+        assert_equal [NoProfileSchema::ExampleExtension], NoProfileSchema::Mutation.all_field_definitions.first.extensions.map(&:class)
+        assert_equal [NoProfileSchema::ExampleExtension], NoProfileSchema::Subscription.all_field_definitions.first.extensions.map(&:class)
+        assert_equal [NoProfileSchema::ExampleExtension, NoProfileSchema::OtherExampleExtension], NoProfileSchema::OrphanType.all_field_definitions.first.extensions.map(&:class)
+        custom_int_field = NoProfileSchema::CustomIntrospection::DynamicFields.all_field_definitions.find { |f| f.original_name == :__hello }
+        assert_equal [], custom_int_field.extensions
+        NoProfileSchema.introspection(NoProfileSchema::CustomIntrospection)
+        assert_equal [NoProfileSchema::OtherExampleExtension], custom_int_field.extensions.map(&:class)
       end
     end
   end