Rework how Query type is defined and customized.

elasticgraph-apollo/apollo_tests_implementation/config/products_schema.rb

@@ -8,7 +8,28 @@
 
 # @private
 module ApolloTestImpl
-  module GraphQLSDLEnumeratorExtension
+  federation_version = ENV["TARGET_APOLLO_FEDERATION_VERSION"]
+
+  # Note: this includes many "manual" schema elements (directives, raw SDL, etc) that the
+  # `elasticgraph-apollo` library will generate on our behalf in the future. For now, this
+  # includes all these schema elements just to make it as close as possible to an apollo
+  # compatible schema without further changes to elasticgraph-apollo.
+  #
+  # https://github.com/apollographql/apollo-federation-subgraph-compatibility/blob/2.0.0/COMPATIBILITY.md#products-schema-to-be-implemented-by-library-maintainers
+  ElasticGraph.define_schema do |schema|
+    schema.json_schema_version 1
+    schema.target_apollo_federation_version(federation_version) if federation_version
+
+    unless federation_version == "2.0"
+      schema.raw_sdl <<~EOS
+        extend schema
+          @link(url: "https://myspecs.dev/myCustomDirective/v1.0", import: ["@custom"])
+          @composeDirective(name: "@custom")
+
+        directive @custom on OBJECT
+      EOS
+    end
+
     # The `apollo-federation-subgraph-compatibility` project requires[^1] that each tested implementation provide
     # specific `Query` fields:
     #
@@ -24,8 +45,8 @@ module GraphQLSDLEnumeratorExtension
     # `Query` type to add the required fields.
     #
     # [^1]: https://github.com/apollographql/apollo-federation-subgraph-compatibility/blob/2.0.0/COMPATIBILITY.md#products-schema-to-be-implemented-by-library-maintainers
-    def root_query_type
-      super.tap do |type|
+    schema.on_built_in_types do |type|
+      if type.name == "Query"
         type.field "product", "Product" do |f|
           f.argument "id", "ID!"
         end
@@ -37,40 +58,6 @@ def root_query_type
         end
       end
     end
-  end
-
-  # @private
-  module SchemaDefFactoryExtension
-    def new_graphql_sdl_enumerator(all_types_except_root_query_type)
-      super(all_types_except_root_query_type).tap do |enum|
-        enum.extend GraphQLSDLEnumeratorExtension
-      end
-    end
-  end
-
-  federation_version = ENV["TARGET_APOLLO_FEDERATION_VERSION"]
-
-  # Note: this includes many "manual" schema elements (directives, raw SDL, etc) that the
-  # `elasticgraph-apollo` library will generate on our behalf in the future. For now, this
-  # includes all these schema elements just to make it as close as possible to an apollo
-  # compatible schema without further changes to elasticgraph-apollo.
-  #
-  # https://github.com/apollographql/apollo-federation-subgraph-compatibility/blob/2.0.0/COMPATIBILITY.md#products-schema-to-be-implemented-by-library-maintainers
-  ElasticGraph.define_schema do |schema|
-    schema.factory.extend SchemaDefFactoryExtension
-
-    schema.json_schema_version 1
-    schema.target_apollo_federation_version(federation_version) if federation_version
-
-    unless federation_version == "2.0"
-      schema.raw_sdl <<~EOS
-        extend schema
-          @link(url: "https://myspecs.dev/myCustomDirective/v1.0", import: ["@custom"])
-          @composeDirective(name: "@custom")
-
-        directive @custom on OBJECT
-      EOS
-    end
 
     schema.object_type "Product" do |t|
       t.directive "custom" unless federation_version == "2.0"

elasticgraph-apollo/lib/elastic_graph/apollo/schema_definition/api_extension.rb

@@ -117,7 +117,12 @@ def self.extended(api)
           api.on_built_in_types do |type|
             # Built-in types like `PageInfo` need to be tagged with `@shareable` on Federation V2 since other subgraphs may
             # have them and they aren't entity types. `Query`, as the root, is a special case that must be skipped.
-            (_ = type).apollo_shareable if type.respond_to?(:apollo_shareable) && type.name != "Query"
+            # We also need to customize it.
+            if type.name == "Query"
+              customize_root_query_type(_ = type)
+            elsif type.respond_to?(:apollo_shareable)
+              (_ = type).apollo_shareable
+            end
           end
 
           api.register_graphql_extension GraphQL::EngineExtension, defined_at: "elastic_graph/apollo/graphql/engine_extension"
@@ -444,6 +449,53 @@ def validate_entity_types_can_all_be_resolved(entity_types)
             raise Errors::SchemaError, unresolvable_field_errors.join("\n#{"-" * 100}\n")
           end
         end
+
+        private_class_method def self.customize_root_query_type(type)
+          if type.schema_def_state.object_types_by_name.values.any?(&:indexed?)
+            type.field "_entities", "[_Entity]!" do |f|
+              f.documentation <<~EOS
+                A field required by the [Apollo Federation subgraph
+                spec](https://www.apollographql.com/docs/federation/subgraph-spec/#query_entities):
+
+                > The graph router uses this root-level `Query` field to directly fetch fields of entities defined by a subgraph.
+                >
+                > This field must take a `representations` argument of type `[_Any!]!` (a non-nullable list of non-nullable
+                > [`_Any` scalars](https://www.apollographql.com/docs/federation/subgraph-spec/#scalar-_any)). Its return type must be `[_Entity]!` (a non-nullable list of _nullable_
+                > objects that belong to the [`_Entity` union](https://www.apollographql.com/docs/federation/subgraph-spec/#union-_entity)).
+                >
+                > Each entry in the `representations` list  must be validated with the following rules:
+                >
+                > - A representation must include a `__typename` string field.
+                > - A representation must contain all fields included in the fieldset of a `@key` directive applied to the corresponding entity definition.
+                >
+                > For details, see [Resolving entity fields with `Query._entities`](https://www.apollographql.com/docs/federation/subgraph-spec/#resolving-entity-fields-with-query_entities).
+
+                Not intended for use by clients other than Apollo.
+              EOS
+
+              f.argument "representations", "[_Any!]!" do |a|
+                a.documentation <<~EOS
+                  A list of entity data blobs from other apollo subgraphs. For more information (and
+                  to see an example of what form this argument takes), see the [Apollo Federation subgraph
+                  spec](https://www.apollographql.com/docs/federation/subgraph-spec/#resolve-requests-for-entities).
+                EOS
+              end
+            end
+          end
+
+          type.field "_service", "_Service!" do |f|
+            f.documentation <<~EOS
+              A field required by the [Apollo Federation subgraph
+              spec](https://www.apollographql.com/docs/federation/subgraph-spec/#query_service):
+
+              > This field of the root `Query` type must return a non-nullable [`_Service` type](https://www.apollographql.com/docs/federation/subgraph-spec/#type-_service).
+
+              > For details, see [Enhanced introspection with `Query._service`](https://www.apollographql.com/docs/federation/subgraph-spec/#enhanced-introspection-with-query_service).
+
+              Not intended for use by clients other than Apollo.
+            EOS
+          end
+        end
       end
     end
   end

elasticgraph-apollo/lib/elastic_graph/apollo/schema_definition/factory_extension.rb

@@ -10,7 +10,6 @@
 require "elastic_graph/apollo/schema_definition/enum_type_extension"
 require "elastic_graph/apollo/schema_definition/enum_value_extension"
 require "elastic_graph/apollo/schema_definition/field_extension"
-require "elastic_graph/apollo/schema_definition/graphql_sdl_enumerator_extension"
 require "elastic_graph/apollo/schema_definition/input_type_extension"
 require "elastic_graph/apollo/schema_definition/interface_type_extension"
 require "elastic_graph/apollo/schema_definition/object_type_extension"
@@ -32,12 +31,6 @@ module FactoryExtension
                      end
                    end
 
-        def new_graphql_sdl_enumerator(all_types_except_root_query_type)
-          super.tap do |enum|
-            enum.extend GraphQLSDLEnumeratorExtension
-          end
-        end
-
         def new_argument(field, name, value_type)
           super(field, name, value_type) do |type|
             type.extend ArgumentExtension

elasticgraph-apollo/sig/elastic_graph/apollo/schema_definition/api_extension.rbs

@@ -19,6 +19,7 @@ module ElasticGraph
         def validate_entity_types_can_all_be_resolved: (::Array[ElasticGraph::SchemaDefinition::indexableType]) -> void
 
         def self.extended: (ElasticGraph::SchemaDefinition::API & APIExtension) -> void
+        def self.customize_root_query_type: (ElasticGraph::SchemaDefinition::SchemaElements::ObjectType) -> void
       end
     end
   end

elasticgraph-graphql/spec/unit/elastic_graph/graphql/query_executor_spec.rb

@@ -361,24 +361,14 @@ class GraphQL
 
         context "when the schema has been customized (as in an extension like elasticgraph-apollo)" do
           before(:context) do
-            enumerator_extension = Module.new do
-              def root_query_type
-                super.tap do |type|
+            self.schema_artifacts = generate_schema_artifacts do |schema|
+              schema.on_built_in_types do |type|
+                if type.name == "Query"
                   type.field "multiply", "Int" do |f|
                     f.argument("operands", "Operands!")
                   end
                 end
               end
-            end
-
-            self.schema_artifacts = generate_schema_artifacts do |schema|
-              schema.factory.extend(Module.new {
-                define_method :new_graphql_sdl_enumerator do |all_types_except_root_query_type|
-                  super(all_types_except_root_query_type).tap do |enum|
-                    enum.extend enumerator_extension
-                  end
-                end
-              })
 
               schema.scalar_type "Operands" do |t|
                 t.mapping type: nil

elasticgraph-schema_definition/lib/elastic_graph/schema_definition/factory.rb

@@ -115,8 +115,8 @@ def new_enums_for_indexed_types
                  end
       @@field_new = prevent_non_factory_instantiation_of(SchemaElements::Field)
 
-      def new_graphql_sdl_enumerator(all_types_except_root_query_type)
-        @@graphql_sdl_enumerator_new.call(@state, all_types_except_root_query_type)
+      def new_graphql_sdl_enumerator(all_types)
+        @@graphql_sdl_enumerator_new.call(@state, all_types)
       end
       @@graphql_sdl_enumerator_new = prevent_non_factory_instantiation_of(SchemaElements::GraphQLSDLEnumerator)