Add JSON ingestion indexing extensions

Author: jwils

elasticgraph-json_ingestion/lib/elastic_graph/json_ingestion/schema_definition/indexing/field_reference_extension.rb

@@ -0,0 +1,61 @@
+# Copyright 2024 - 2026 Block, Inc.
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+#
+# frozen_string_literal: true
+
+require "elastic_graph/json_ingestion/schema_definition/indexing/field"
+
+module ElasticGraph
+  module JSONIngestion
+    module SchemaDefinition
+      module Indexing
+        # Extends indexing field references with JSON schema state needed when resolving fields.
+        #
+        # @api private
+        module FieldReferenceExtension
+          def with_json_schema(json_schema_layers:, json_schema_customizations:)
+            @json_schema_layers = json_schema_layers
+            @json_schema_customizations = json_schema_customizations
+            self
+          end
+
+          def ==(other)
+            other.is_a?(FieldReferenceExtension) &&
+              field_reference_comparison_metadata == other.field_reference_comparison_metadata &&
+              json_schema_comparison_metadata == other.json_schema_comparison_metadata
+          end
+
+          def resolve
+            return nil unless (resolved_field = super)
+
+            json_schema_field = resolved_field.extend(Indexing::FieldExtension) # : ElasticGraph::SchemaDefinition::Indexing::Field & FieldExtension
+            json_schema_field.with_json_schema(
+              json_schema_layers: @json_schema_layers,
+              json_schema_customizations: @json_schema_customizations
+            )
+          end
+
+          def json_schema_comparison_metadata
+            [@json_schema_layers, @json_schema_customizations]
+          end
+
+          def field_reference_comparison_metadata
+            [
+              name,
+              name_in_index,
+              type,
+              mapping_options,
+              accuracy_confidence,
+              source,
+              runtime_field_script,
+              doc_comment
+            ]
+          end
+        end
+      end
+    end
+  end
+end

elasticgraph-json_ingestion/lib/elastic_graph/json_ingestion/schema_definition/indexing/field_type/enum_extension.rb

@@ -0,0 +1,52 @@
+# Copyright 2024 - 2026 Block, Inc.
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+#
+# frozen_string_literal: true
+
+module ElasticGraph
+  module JSONIngestion
+    module SchemaDefinition
+      module Indexing
+        # Namespace for indexing-field-type extensions that contribute JSON schema generation behavior.
+        module FieldType
+          # Extends enum indexing field types with JSON schema serialization.
+          #
+          # @private
+          module EnumExtension
+            # @return [Hash<String, ::Object>] additional ElasticGraph metadata to put in the JSON schema for this enum type.
+            def json_schema_field_metadata_by_field_name
+              {}
+            end
+
+            # @param customizations [Hash<String, ::Object>] JSON schema customizations
+            # @return [Hash<String, ::Object>] formatted customizations.
+            def format_field_json_schema_customizations(customizations)
+              # Since an enum type already restricts the values to a small set of allowed values, we do not need to keep
+              # other customizations (such as the `maxLength` field customization EG automatically applies to fields
+              # indexed as a `keyword`--we don't allow enum values to exceed that length, anyway).
+              #
+              # It's desirable to restrict what customizations are applied because when a publisher uses the JSON schema
+              # to generate code using a library such as https://github.com/pwall567/json-kotlin-schema-codegen, we found
+              # that the presence of extra field customizations inhibits the library's ability to generate code in the way
+              # we want (it causes the type of the enum to change since the JSON schema changes from a direct `$ref` to
+              # being wrapped in an `allOf`).
+              #
+              # However, we still want to apply `enum` customizations--this allows a user to "narrow" the set of allowed
+              # values for a field. For example, a `Currency` enum could contain every currency, and a user may want to
+              # restrict a specific `currency` field to a subset of currencies (e.g. to just USD, CAD, and EUR).
+              customizations.slice("enum")
+            end
+
+            # @return [Hash<String, ::Object>] the JSON schema for this enum type.
+            def to_json_schema
+              {"type" => "string", "enum" => enum_value_names}
+            end
+          end
+        end
+      end
+    end
+  end
+end

elasticgraph-json_ingestion/lib/elastic_graph/json_ingestion/schema_definition/indexing/field_type/object_extension.rb

@@ -0,0 +1,101 @@
+# Copyright 2024 - 2026 Block, Inc.
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+#
+# frozen_string_literal: true
+
+require "elastic_graph/support/hash_util"
+
+module ElasticGraph
+  module JSONIngestion
+    module SchemaDefinition
+      module Indexing
+        module FieldType
+          # Extends object/interface indexing field types with JSON schema serialization.
+          #
+          # @private
+          module ObjectExtension
+            def json_schema_options
+              @json_schema_options ||= {}
+            end
+
+            def with_json_schema_options(json_schema_options)
+              @json_schema_options = json_schema_options
+              self
+            end
+
+            # @return [Hash<String, JSONSchemaFieldMetadata>] field metadata keyed by field name
+            def json_schema_field_metadata_by_field_name
+              subfields.to_h { |field| [field.name, field.json_schema_metadata] }
+            end
+
+            # @param customizations [Hash<String, Object>] the customizations to format
+            # @return [Hash<String, Object>] the formatted customizations
+            def format_field_json_schema_customizations(customizations)
+              customizations
+            end
+
+            # @return [Hash<String, Object>] the JSON schema definition for this object type
+            def to_json_schema
+              @to_json_schema ||=
+                if json_schema_options.empty?
+                  # Fields that are `sourced_from` an alternate type must not be included in this type's JSON schema,
+                  # since events of this type won't include them.
+                  other_source_subfields, json_schema_candidate_subfields = subfields.partition(&:source)
+                  validate_sourced_fields_have_no_json_schema_overrides(other_source_subfields)
+                  json_schema_subfields = json_schema_candidate_subfields.reject(&:runtime_field_script)
+                  required_fields = json_schema_subfields
+                  required_fields = required_fields.reject(&:nullable?) if schema_def_state.allow_omitted_json_schema_fields
+
+                  {
+                    "type" => "object",
+                    "properties" => json_schema_subfields.to_h { |field| [field.name, field.json_schema] }.merge(json_schema_typename_field),
+                    # Note: `__typename` is intentionally not included in the `required` list. If `__typename` is present
+                    # we want it validated (as we do by merging in `json_schema_typename_field`) but we only want
+                    # to require it in the context of a union type. The union's JSON schema requires the field.
+                    "required" => required_fields.map(&:name).freeze,
+                    "additionalProperties" => (false unless schema_def_state.allow_extra_json_schema_fields),
+                    "description" => doc_comment
+                  }.compact.freeze
+                else
+                  Support::HashUtil.stringify_keys(json_schema_options)
+                end
+            end
+
+            private
+
+            # Returns a `__typename` property which we use for union types.
+            #
+            # This must always be set to the name of the type (thus the const value).
+            #
+            # We also add a "default" value. This does not impact validation, but rather
+            # aids tools like our Kotlin codegen to save publishers from having to set the
+            # property explicitly when creating events.
+            def json_schema_typename_field
+              {
+                "__typename" => {
+                  "type" => "string",
+                  "const" => type_name,
+                  "default" => type_name
+                }
+              }
+            end
+
+            def validate_sourced_fields_have_no_json_schema_overrides(other_source_subfields)
+              problem_fields = other_source_subfields.reject { |field| field.json_schema_customizations.empty? }
+              return if problem_fields.empty?
+
+              field_descriptions = problem_fields.map(&:name).sort.map { |field| "`#{field}`" }.join(", ")
+              raise Errors::SchemaError,
+                "`#{type_name}` has #{problem_fields.size} field(s) (#{field_descriptions}) that are `sourced_from` " \
+                "another type and also have JSON schema customizations. Instead, put the JSON schema " \
+                "customizations on the source type's field definitions."
+            end
+          end
+        end
+      end
+    end
+  end
+end

elasticgraph-json_ingestion/lib/elastic_graph/json_ingestion/schema_definition/indexing/field_type/scalar_extension.rb

@@ -0,0 +1,43 @@
+# Copyright 2024 - 2026 Block, Inc.
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+#
+# frozen_string_literal: true
+
+require "elastic_graph/support/hash_util"
+
+module ElasticGraph
+  module JSONIngestion
+    module SchemaDefinition
+      module Indexing
+        module FieldType
+          # Extends scalar indexing field types with JSON schema serialization.
+          #
+          # @private
+          module ScalarExtension
+            # @return [Hash] empty hash, as scalar types have no subfields
+            def json_schema_field_metadata_by_field_name
+              {}
+            end
+
+            # @param customizations [Hash<String, Object>] the customizations to format
+            # @return [Hash<String, Object>] the formatted customizations
+            def format_field_json_schema_customizations(customizations)
+              customizations
+            end
+
+            # @return [Hash<String, Object>] the JSON schema definition for this scalar type
+            def to_json_schema
+              json_scalar_type = scalar_type # : ElasticGraph::SchemaDefinition::SchemaElements::ScalarType & SchemaElements::ScalarTypeExtension
+              json_scalar_type.validate_json_schema_configuration!
+
+              Support::HashUtil.stringify_keys(json_scalar_type.json_schema_options)
+            end
+          end
+        end
+      end
+    end
+  end
+end

elasticgraph-json_ingestion/lib/elastic_graph/json_ingestion/schema_definition/indexing/field_type/union_extension.rb

@@ -0,0 +1,52 @@
+# Copyright 2024 - 2026 Block, Inc.
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+#
+# frozen_string_literal: true
+
+module ElasticGraph
+  module JSONIngestion
+    module SchemaDefinition
+      module Indexing
+        module FieldType
+          # Extends union indexing field types with JSON schema serialization.
+          #
+          # @private
+          module UnionExtension
+            # @return [Hash] empty hash, as union types have no subfields
+            def json_schema_field_metadata_by_field_name
+              {}
+            end
+
+            # @param customizations [Hash<String, Object>] the customizations to format
+            # @return [Hash<String, Object>] the formatted customizations
+            def format_field_json_schema_customizations(customizations)
+              customizations
+            end
+
+            # @return [Hash<String, Object>] the JSON schema definition for this union type
+            def to_json_schema
+              subtype_json_schemas = subtypes_by_name.keys.map { |name| {"$ref" => "#/$defs/#{name}"} }
+
+              # A union type can represent multiple subtypes, referenced by the "anyOf" clause below.
+              # We also add a requirement for the presence of __typename to indicate which type
+              # is being referenced (this property is pre-defined on the type itself as a constant).
+              #
+              # Note: Although both "oneOf" and "anyOf" keywords are valid for combining schemas
+              # to form a union, and validate equivalently when no object can satisfy multiple of the
+              # subschemas (which is the case here given the __typename requirements are mutually
+              # exclusive), we chose to use "oneOf" here because it works better with this library:
+              # https://github.com/pwall567/json-kotlin-schema-codegen
+              {
+                "required" => %w[__typename],
+                "oneOf" => subtype_json_schemas
+              }
+            end
+          end
+        end
+      end
+    end
+  end
+end

elasticgraph-json_ingestion/sig/elastic_graph/json_ingestion/schema_definition/indexing/field_reference_extension.rbs

@@ -0,0 +1,34 @@
+module ElasticGraph
+  module JSONIngestion
+    module SchemaDefinition
+      module Indexing
+        module FieldReferenceExtension: ::ElasticGraph::SchemaDefinition::Indexing::FieldReference
+          @json_schema_layers: ::ElasticGraph::SchemaDefinition::jsonSchemaLayersArray
+          @json_schema_customizations: ::Hash[::Symbol, untyped]
+
+          def with_json_schema: (
+            json_schema_layers: ::ElasticGraph::SchemaDefinition::jsonSchemaLayersArray,
+            json_schema_customizations: ::Hash[::Symbol, untyped]
+          ) -> (::ElasticGraph::SchemaDefinition::Indexing::FieldReference & FieldReferenceExtension)
+          def ==: (untyped) -> bool
+          def resolve: () -> (::ElasticGraph::SchemaDefinition::Indexing::Field & FieldExtension)?
+
+          def json_schema_comparison_metadata: () -> [
+            ::ElasticGraph::SchemaDefinition::jsonSchemaLayersArray,
+            ::Hash[::Symbol, untyped]
+          ]
+          def field_reference_comparison_metadata: () -> [
+            ::String,
+            ::String,
+            ::ElasticGraph::SchemaDefinition::SchemaElements::TypeReference,
+            ::Hash[::Symbol, untyped],
+            ::ElasticGraph::SchemaDefinition::Indexing::Field::accuracyConfidence,
+            ::ElasticGraph::SchemaDefinition::SchemaElements::FieldSource?,
+            ::String?,
+            ::String?
+          ]
+        end
+      end
+    end
+  end
+end

elasticgraph-json_ingestion/sig/elastic_graph/json_ingestion/schema_definition/indexing/field_type.rbs

@@ -0,0 +1,34 @@
+module ElasticGraph
+  module JSONIngestion
+    module SchemaDefinition
+      module Indexing
+        interface _EnumFieldType
+          def enum_value_names: () -> ::Array[::String]
+        end
+
+        interface _JSONFieldType
+          def to_mapping: () -> ::Hash[::String, untyped]
+          def json_schema_field_metadata_by_field_name: () -> ::Hash[::String, JSONSchemaFieldMetadata]
+          def format_field_json_schema_customizations: (::Hash[::String, untyped]) -> ::Hash[::String, untyped]
+          def to_json_schema: () -> ::Hash[::String, untyped]
+        end
+
+        interface _ObjectFieldType
+          def schema_def_state: () -> ::ElasticGraph::SchemaDefinition::State
+          def type_name: () -> ::String
+          def subfields: () -> ::Array[::ElasticGraph::SchemaDefinition::Indexing::Field]
+          def mapping_options: () -> ::ElasticGraph::SchemaDefinition::Mixins::HasTypeInfo::optionsHash
+          def doc_comment: () -> ::String?
+        end
+
+        interface _ScalarFieldType
+          def scalar_type: () -> ::ElasticGraph::SchemaDefinition::SchemaElements::ScalarType
+        end
+
+        interface _UnionFieldType
+          def subtypes_by_name: () -> ::Hash[::String, ::Object]
+        end
+      end
+    end
+  end
+end

elasticgraph-json_ingestion/sig/elastic_graph/json_ingestion/schema_definition/indexing/field_type/enum_extension.rbs

@@ -0,0 +1,15 @@
+module ElasticGraph
+  module JSONIngestion
+    module SchemaDefinition
+      module Indexing
+        module FieldType
+          module EnumExtension: _EnumFieldType
+            def json_schema_field_metadata_by_field_name: () -> ::Hash[::String, JSONSchemaFieldMetadata]
+            def format_field_json_schema_customizations: (::Hash[::String, untyped]) -> ::Hash[::String, untyped]
+            def to_json_schema: () -> ::Hash[::String, untyped]
+          end
+        end
+      end
+    end
+  end
+end