Add JSON ingestion indexing extensions

Author: jwils

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

@@ -0,0 +1,136 @@
+# 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/constants"
+require "elastic_graph/json_ingestion/schema_definition/indexing/json_schema_field_metadata"
+require "elastic_graph/support/hash_util"
+
+module ElasticGraph
+  module JSONIngestion
+    module SchemaDefinition
+      # Namespace for JSON-schema-aware indexing components.
+      module Indexing
+        # Extends indexing fields with JSON schema generation behavior.
+        #
+        # @api private
+        module FieldExtension
+          # JSON schema overrides that automatically apply to specific mapping types so that the JSON schema
+          # validation will reject values which cannot be indexed into fields of a specific mapping type.
+          #
+          # @see https://www.elastic.co/guide/en/elasticsearch/reference/current/number.html Elasticsearch numeric field type documentation
+          # @note We don't handle `integer` here because it's the default numeric type (handled by our definition of the `Int` scalar type).
+          # @note Likewise, we don't handle `long` here because a custom scalar type must be used for that since GraphQL's `Int` type can't handle long values.
+          JSON_SCHEMA_OVERRIDES_BY_MAPPING_TYPE = {
+            "byte" => {"minimum" => -(2**7), "maximum" => (2**7) - 1},
+            "short" => {"minimum" => -(2**15), "maximum" => (2**15) - 1},
+            "keyword" => {"maxLength" => DEFAULT_MAX_KEYWORD_LENGTH},
+            "text" => {"maxLength" => DEFAULT_MAX_TEXT_LENGTH}
+          }
+
+          # @return [Hash<Symbol, Object>] user-specified JSON schema customizations for this field
+          def json_schema_customizations
+            @json_schema_customizations
+          end
+
+          # @private
+          def with_json_schema(json_schema_layers:, json_schema_customizations:)
+            @json_schema_layers = json_schema_layers
+            @json_schema_customizations = json_schema_customizations
+            self
+          end
+
+          # Returns the JSON schema definition for this field.
+          #
+          # @return [Hash<String, Object>] the JSON schema hash
+          def json_schema
+            @json_schema ||=
+              json_schema_layers
+                .reverse # resolve layers from innermost to outermost wrappings
+                .reduce(inner_json_schema) { |acc, layer| process_layer(layer, acc) }
+                .merge(outer_json_schema_customizations)
+                .merge({"description" => doc_comment}.compact)
+                .then { |hash| Support::HashUtil.stringify_keys(hash) }
+          end
+
+          # @return [JSONSchemaFieldMetadata] metadata about this field for inclusion in the JSON schema
+          def json_schema_metadata
+            JSONSchemaFieldMetadata.new(type: type.name, name_in_index: name_in_index)
+          end
+
+          def nullable?
+            json_schema_layers.include?(:nullable)
+          end
+
+          private
+
+          def json_schema_layers
+            @json_schema_layers
+          end
+
+          def inner_json_schema
+            user_specified_customizations =
+              if user_specified_json_schema_customizations_go_on_outside?
+                {} # : ::Hash[::String, untyped]
+              else
+                Support::HashUtil.stringify_keys(json_schema_customizations)
+              end
+
+            customizations_from_mapping = JSON_SCHEMA_OVERRIDES_BY_MAPPING_TYPE[mapping["type"]] || {}
+            customizations = customizations_from_mapping.merge(user_specified_customizations)
+            # @type var field_type: _JSONFieldType
+            field_type = _ = indexing_field_type
+            customizations = field_type.format_field_json_schema_customizations(customizations)
+
+            ref = {"$ref" => "#/$defs/#{type.unwrapped_name}"}
+            return ref if customizations.empty?
+
+            # Combine any customizations with the type ref under an "allOf" subschema:
+            # all of these properties must hold true for the type to be valid.
+            #
+            # Note that if we simply combine the customizations with the `$ref`
+            # at the same level, it will not work, because other subschema
+            # properties are ignored when they are in the same object as a `$ref`:
+            # https://github.com/json-schema-org/JSON-Schema-Test-Suite/blob/2.0.0/tests/draft7/ref.json#L165-L168
+            {"allOf" => [ref, customizations]}
+          end
+
+          def outer_json_schema_customizations
+            return {} unless user_specified_json_schema_customizations_go_on_outside?
+            Support::HashUtil.stringify_keys(json_schema_customizations)
+          end
+
+          # Indicates if the user-specified JSON schema customizations should go on the inside
+          # (where they normally go) or on the outside. They only go on the outside when it's
+          # an array field, because then they apply to the array itself instead of the items in the
+          # array.
+          def user_specified_json_schema_customizations_go_on_outside?
+            json_schema_layers.include?(:array)
+          end
+
+          def process_layer(layer, schema)
+            case layer
+            when :nullable
+              # Here we use "anyOf" to ensure that JSON can either match the schema OR null.
+              #
+              # (Using "oneOf" would mean that if we had a schema that also allowed null,
+              # null would never be allowed, since "oneOf" must match exactly one subschema).
+              {
+                "anyOf" => [
+                  schema,
+                  {"type" => "null"}
+                ]
+              }
+            when :array
+              {"type" => "array", "items" => schema}
+            end
+          end
+        end
+      end
+    end
+  end
+end

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,40 @@
+# 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
+              Support::HashUtil.stringify_keys(scalar_type.json_schema_options)
+            end
+          end
+        end
+      end
+    end
+  end
+end