Schema support for inherited indexes

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

@@ -336,6 +336,13 @@ def define_apollo_schema_elements
             end
           end
 
+          # Add @key directives to all root document types with an id field.
+          # This happens after schema definition is complete so root_document_type? sees the complete type hierarchy.
+          state.object_types_by_name.values
+            .grep(ElasticGraph::SchemaDefinition::SchemaElements::ObjectType)
+            .select { |object_type| object_type.root_document_type? && object_type.graphql_fields_by_name.key?("id") }
+            .each { |object_type| object_type.apollo_key fields: "id" }
+
           entity_types = state.object_types_by_name.values.select do |object_type|
             object_type.directives.any? do |directive|
               directive.name == "key" && directive.arguments.fetch(:resolvable, true)

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

@@ -9,23 +9,29 @@
 module ElasticGraph
   module Apollo
     module SchemaDefinition
-      # The Apollo `_Entity` type is a type union of _all_ entity subtypes in an ElasticGraph schema.
-      # However, unlike a normal union type:
-      #
-      # - `_Entity` is never a root document type, and should not be treated as one (even though its subtypes are all root document types,
-      #   which would usually cause it to be treated as a root document type!).
-      # - A merged set of `graphql_fields_by_name` cannot be safely computed. That method raises errors if a field with the same name
-      #   has conflicting definitions on different subtypes, but we must allow that on `_Entity` subtypes.
+      # The Apollo `_Entity` type is a union of all entity types in an ElasticGraph schema. These overrides
+      # prevent ElasticGraph from treating `_Entity` like a normal indexed union type, which would trigger
+      # unwanted derived schema generation and validation.
       #
       # @private
       module EntityTypeExtension
+        # A merged set of `graphql_fields_by_name` cannot be safely computed. That method raises errors if a field with
+        # the same name has conflicting definitions on different subtypes, but we must allow that on `_Entity` subtypes.
         def graphql_fields_by_name
           {}
         end
 
+        # `_Entity` is never a root document type, and should not be treated as one (even though its subtypes are all
+        # root document types, which would usually cause it to be treated as a root document type!).
         def root_document_type?
           false
         end
+
+        # `_Entity` is never directly queryable from the root `Query` type. It's queried via the apollo
+        # `_entities(representations: ...)` field instead.
+        def directly_queryable?
+          false
+        end
       end
     end
   end

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

@@ -66,17 +66,12 @@ def new_interface_type(name)
           end
         end
 
-        # Here we override `object_type` in order to automatically add the apollo `@key` directive to indexed types.
         def new_object_type(name)
           super(name) do |raw_type|
             raw_type.extend ObjectTypeExtension
             type = raw_type # : ElasticGraph::SchemaDefinition::SchemaElements::ObjectType & ObjectTypeExtension
 
             yield type if block_given?
-
-            if type.root_document_type? && type.graphql_fields_by_name.key?("id")
-              type.apollo_key fields: "id"
-            end
           end
         end
 

elasticgraph-apollo/spec/unit/elastic_graph/apollo/schema_definition_spec.rb

@@ -291,24 +291,27 @@ def self.with_both_casing_forms(&block)
 
         it "avoids including indexed interfaces in the `_Entity` union (and does not add `@key` to it) since unions can't include interfaces" do
           schema_string = graphql_schema_string do |schema|
+            # Define subtypes before their indexed supertype to verify type references
+            # resolve correctly regardless of definition order.
             schema.object_type "IndexedType1" do |t|
               t.implements "NamedEntity"
               t.field "graphql", "String", name_in_index: "index"
               t.field "id", "ID!"
               t.field "name", "String"
-              t.index "index1"
+              # Inherits index from NamedEntity
             end
 
             schema.object_type "IndexedType2" do |t|
               t.implements "NamedEntity"
               t.field "id", "ID!"
               t.field "name", "String"
-              t.index "index1"
+              # Inherits index from NamedEntity
             end
 
             schema.interface_type "NamedEntity" do |t|
               t.field "id", "ID!"
               t.field "name", "String"
+              t.index "named_entities"
             end
           end
 

elasticgraph-graphql/spec/unit/elastic_graph/graphql/schema/type_spec.rb

@@ -53,12 +53,22 @@ class Schema
                 t.field "name", "String"
               end
 
+              schema.object_type "Temperature" do |t|
+                t.field "id", "ID"
+                t.field "value", "Float"
+              end
+
+              schema.object_type "Pressure" do |t|
+                t.field "id", "ID"
+                t.field "amount", "Int"
+              end
+
               schema.union_type "Attribute" do |t|
-                t.subtypes "Color", "Velocity"
+                t.subtypes "Color"
               end
 
               schema.union_type "IndexedAttribute" do |t|
-                t.subtypes "Color", "Velocity"
+                t.subtypes "Temperature", "Pressure"
                 t.index "attributes"
               end
 
@@ -266,6 +276,26 @@ class Schema
             expect(type.unwrap_non_null).to be type
           end
 
+          it "can model a type that inherits an index from a union" do
+            # Temperature doesn't have its own index, but inherits from IndexedAttribute union
+            type = schema.type_named("Temperature")
+
+            expect(type.name).to eq "Temperature"
+            expect(type).to only_satisfy_predicates(:nullable?, :object?, :indexed_document?)
+            expect(type.unwrap_fully).to be schema.type_named("Temperature")
+            expect(type.unwrap_non_null).to be type
+          end
+
+          it "can model a type that inherits an index from an interface" do
+            # Velocity doesn't have its own index, but inherits from DirectlyIndexedInterface
+            type = schema.type_named("Velocity")
+
+            expect(type.name).to eq "Velocity"
+            expect(type).to only_satisfy_predicates(:nullable?, :object?, :indexed_document?)
+            expect(type.unwrap_fully).to be schema.type_named("Velocity")
+            expect(type.unwrap_non_null).to be type
+          end
+
           it "can model an indexed aggregation type" do
             type = type_for("indexed_aggregation")
 

elasticgraph-schema_definition/lib/elastic_graph/schema_definition/mixins/has_indices.rb

@@ -100,9 +100,39 @@ def has_own_index_def?
           !@own_index_def.nil?
         end
 
-        # @return [Boolean] true if this type is a root document type that lives at the document root in the datastore.
-        #   For types with `own_index_def`, returns true. For abstract types with indexed subtypes, overridden in {HasSubtypes}.
+        # Resolves this type's index definition. This will be one of:
+        # - This type's own_index_def (if it directly defines an index)
+        # - An inherited index from an abstract supertype (union/interface) that has an index
+        #
+        # This type can be a subtype of multiple abstract types (e.g., implements multiple interfaces), but unless it
+        # defines its own index, at most one of its supertypes may have an index. If multiple parent types are indexed,
+        # this method raises an error to prevent ambiguity about which index to inherit.
+        #
+        # @return [Indexing::Index, nil] the index definition, or nil if this type has no index
+        # @raise [Errors::SchemaError] if this type is a subtype of multiple indexed abstract types
+        def index_def
+          return own_index_def if has_own_index_def?
+
+          indexed_supertypes = recursively_resolve_supertypes.select(&:has_own_index_def?)
+
+          if indexed_supertypes.size > 1
+            parent_names = indexed_supertypes.map { |p| p.own_index_def.name }.join(", ")
+            raise Errors::SchemaError,
+              "The `#{name}` type is a subtype of multiple indexed abstract types (#{parent_names}). " \
+              "If a concrete type does not define an index, it may not be a member of multiple indexed abstract types."
+          end
+
+          indexed_supertypes.first&.own_index_def
+        end
+
+        # @return [Boolean] true if this type is a root document type that lives at a document root in the datastore (is indexed).
+        #   This returns true for types with their own index definition or types that inherit an index from a supertype.
         def root_document_type?
+          !index_def.nil?
+        end
+
+        # @return [Boolean] true if this type is directly queryable via a type-specific field on the root `Query` type.
+        def directly_queryable?
           has_own_index_def?
         end
 
@@ -191,7 +221,7 @@ def override_runtime_metadata(**overrides)
         def runtime_metadata(extra_update_targets)
           SchemaArtifacts::RuntimeMetadata::ObjectType.new(
             update_targets: derived_indexed_types.map(&:runtime_metadata_for_source_type) + [self_update_target].compact + extra_update_targets,
-            index_definition_names: [own_index_def&.name].compact,
+            index_definition_names: [index_def&.name].compact,
             graphql_fields_by_name: runtime_metadata_graphql_fields_by_name,
             elasticgraph_category: nil,
             source_type: nil,
@@ -284,7 +314,7 @@ def self_update_target
             [field, SchemaArtifacts::RuntimeMetadata::DynamicParam.new(source_path: field, cardinality: :one)]
           end
 
-          index_runtime_metadata = own_index_def.runtime_metadata
+          index_runtime_metadata = index_def.runtime_metadata
 
           Indexing::UpdateTargetFactory.new_normal_indexing_update_target(
             type: name,

elasticgraph-schema_definition/lib/elastic_graph/schema_definition/mixins/has_subtypes.rb

@@ -38,6 +38,13 @@ def root_document_type?
           super || subtypes_are_root_document_types?
         end
 
+        # An abstract type is queryable if all of its subtypes are root document types (via a direct or inherited index)
+        # even if those subtypes aren't themselves directly queryable. This is why this doesn't delegate to a
+        # subtypes_are_directly_queryable helper.
+        def directly_queryable?
+          super || subtypes_are_root_document_types?
+        end
+
         def recursively_resolve_subtypes
           resolve_subtypes.flat_map do |type|
             type.is_a?(HasSubtypes) ? (_ = type).recursively_resolve_subtypes : [type]

elasticgraph-schema_definition/lib/elastic_graph/schema_definition/mixins/implements_interfaces.rb

@@ -117,6 +117,21 @@ def to_sdl(&field_arg_selector)
 
           generate_sdl(name_section: name_section, &field_arg_selector)
         end
+
+        # Returns all supertypes of this type, including union memberships and interface ancestors.
+        #
+        # @return [Set<UnionType, InterfaceType>] set of supertypes
+        # @private
+        def recursively_resolve_supertypes
+          union_memberships = schema_def_state.union_types_by_member_ref[type_ref]
+
+          interface_supertypes = implemented_interfaces.flat_map do |interface_ref|
+            interface = interface_ref.resolved
+            [interface] + interface.recursively_resolve_supertypes.to_a
+          end.to_set
+
+          union_memberships | interface_supertypes
+        end
       end
     end
   end

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

@@ -119,7 +119,7 @@ def define_root_graphql_type
           query_type.documentation "The query entry point for the entire schema."
           query_type.resolve_fields_with nil
 
-          state.types_by_name.values.select(&:root_document_type?).sort_by(&:name).each do |type|
+          state.object_types_by_name.values.select(&:directly_queryable?).sort_by(&:name).each do |type|
             # @type var root_doc_type: Mixins::HasIndices & _Type
             root_doc_type = _ = type
 
@@ -212,7 +212,7 @@ def build_runtime_metadata
         enum_generator = state.factory.new_enums_for_root_document_types
 
         sort_order_enum_types_by_name = state.object_types_by_name.values
-          .select(&:root_document_type?)
+          .select(&:directly_queryable?)
           .filter_map { |type| enum_generator.sort_order_enum_for(_ = type) }
           .to_h { |enum_type| [(_ = enum_type).name, (_ = enum_type).runtime_metadata] }
 

elasticgraph-schema_definition/lib/elastic_graph/schema_definition/schema_elements/enum_type.rb

@@ -161,11 +161,6 @@ def to_indexing_field_type
           Indexing::FieldType::Enum.new(values_by_name.keys)
         end
 
-        # @return [false] enum types are never root document types
-        def root_document_type?
-          false
-        end
-
         # @return [EnumType] converts the enum type to its input form for when different naming is used for input vs output enums.
         def as_input
           input_name = type_ref

elasticgraph-schema_definition/lib/elastic_graph/schema_definition/schema_elements/object_type.rb

@@ -38,6 +38,16 @@ class ObjectType < DelegateClass(TypeWithSubfields)
         include Mixins::ImplementsInterfaces
         include Mixins::HasReadableToSAndInspect.new { |t| t.name }
 
+        # @return [Hash<String, Field>] fields that will be indexed, including __typename for mixed-type indices (types
+        # that inherit an index from an abstract supertype)
+        # @private
+        def indexing_fields_by_name_in_index
+          return super if has_own_index_def?
+          return super unless root_document_type?
+
+          super.merge("__typename" => schema_def_state.factory.new_field(name: "__typename", type: "String", parent_type: self))
+        end
+
         # @private
         def initialize(schema_def_state, name)
           field_factory = schema_def_state.factory.method(:new_field)

elasticgraph-schema_definition/lib/elastic_graph/schema_definition/schema_elements/scalar_type.rb

@@ -236,11 +236,6 @@ def derived_graphql_types
           end
         end
 
-        # @private
-        def root_document_type?
-          false
-        end
-
         private
 
         EQUAL_TO_ANY_OF_DOC = <<~EOS

elasticgraph-schema_definition/lib/elastic_graph/schema_definition/schema_elements/type_with_subfields.rb

@@ -470,11 +470,6 @@ def aggregated_values_type
           schema_def_state.type_ref("NonNumeric").as_aggregated_values
         end
 
-        # @private
-        def root_document_type?
-          false
-        end
-
         # @private
         def to_indexing_field_type
           Indexing::FieldType::Object.new(

elasticgraph-schema_definition/lib/elastic_graph/schema_definition/schema_elements/union_type.rb

@@ -94,6 +94,9 @@ def subtype(name)
           end
 
           subtype_refs << type_ref
+
+          # Register reverse lookup so we can efficiently find which unions contain this type
+          schema_def_state.union_types_by_member_ref[type_ref] << self
         end
 
         # Defines multiple subtypes of this union type.
@@ -128,6 +131,15 @@ def to_sdl
           "#{formatted_documentation}union #{name} #{directives_sdl(suffix_with: " ")}= #{subtype_refs.map(&:name).to_a.join(" | ")}"
         end
 
+        # Union types cannot themselves be members of other unions or implement interfaces,
+        # so they have no supertypes.
+        #
+        # @return [Set] empty set
+        # @private
+        def recursively_resolve_supertypes
+          Set[]
+        end
+
         # @private
         def verify_graphql_correctness!
           # Nothing to verify. `verify_graphql_correctness!` will be called on each subtype automatically.

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

@@ -32,6 +32,7 @@ class State < Struct.new(
       :scalar_types_by_name,
       :enum_types_by_name,
       :implementations_by_interface_ref,
+      :union_types_by_member_ref,
       :sdl_parts,
       :paginated_collection_element_types,
       :user_defined_fields,
@@ -79,6 +80,7 @@ def self.with(
           scalar_types_by_name: {},
           enum_types_by_name: {},
           implementations_by_interface_ref: ::Hash.new { |h, k| h[k] = ::Set.new },
+          union_types_by_member_ref: ::Hash.new { |h, k| h[k] = ::Set.new },
           sdl_parts: [],
           paginated_collection_element_types: ::Set.new,
           user_defined_fields: ::Set.new,

elasticgraph-schema_definition/sig/elastic_graph/schema_definition/mixins/has_indices.rbs

@@ -4,7 +4,9 @@ module ElasticGraph
       module HasIndices
         def own_index_def: () -> Indexing::Index?
         def has_own_index_def?: () -> bool
+        def index_def: () -> Indexing::Index?
         def root_document_type?: () -> bool
+        def directly_queryable?: () -> bool
         attr_reader runtime_metadata_overrides: ::Hash[::Symbol, untyped]
         attr_reader default_graphql_resolver: SchemaArtifacts::RuntimeMetadata::ConfiguredGraphQLResolver?
         def index: (::String, ::Hash[::Symbol, ::String | ::Integer]) ?{ (Indexing::Index) -> void } -> Indexing::Index

elasticgraph-schema_definition/sig/elastic_graph/schema_definition/mixins/implements_interfaces.rbs

@@ -5,6 +5,7 @@ module ElasticGraph
         def implements: (*::String) -> void
         attr_reader implemented_interfaces: ::Array[SchemaElements::TypeReference]
         def verify_graphql_correctness!: () -> void
+        def recursively_resolve_supertypes: () -> ::Set[SchemaElements::UnionType | SchemaElements::InterfaceType]
       end
     end
   end