elasticgraph-graphql 0.18.0.0

data/lib/elastic_graph/graphql/aggregation/query.rb

@@ -0,0 +1,172 @@
+# Copyright 2024 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/graphql/aggregation/key"
+require "elastic_graph/graphql/datastore_query"
+require "elastic_graph/graphql/filtering/field_path"
+
+module ElasticGraph
+  class GraphQL
+    module Aggregation
+      class Query < ::Data.define(
+        # Unique name for the aggregation
+        :name,
+        # Whether or not we need to get the document count for each bucket.
+        :needs_doc_count,
+        # Whether or not we need to get the error on the document count to satisfy the sub-aggregation query.
+        # https://www.elastic.co/guide/en/elasticsearch/reference/8.10/search-aggregations-bucket-terms-aggregation.html#_per_bucket_document_count_error
+        :needs_doc_count_error,
+        # Filter to apply to this sub-aggregation.
+        :filter,
+        # Paginator for handling size and other pagination concerns.
+        :paginator,
+        # A sub-aggregation query can have sub-aggregations of its own.
+        :sub_aggregations,
+        # Collection of `Computation` objects that specify numeric computations to perform.
+        :computations,
+        # Collection of `DateHistogramGrouping`, `FieldTermGrouping`, and `ScriptTermGrouping` objects that specify how this sub-aggregation should be grouped.
+        :groupings,
+        # Adapter to use for groupings.
+        :grouping_adapter
+      )
+        def needs_total_doc_count?
+          # We only need a total document count when there are NO groupings and the doc count is requested.
+          # The datastore will return the number of hits in each grouping automatically, so we don't need
+          # a total doc count when there are groupings. And when the query isn't requesting the field, we
+          # don't need it, either.
+          needs_doc_count && groupings.empty?
+        end
+
+        # Builds an aggregations hash. The returned value has a few different cases:
+        #
+        # - If `size` is 0, or `groupings` and `computations` are both empty, we return an empty hash,
+        #   so that `to_datastore_body` is an empty hash. We do this so that we avoid sending
+        #   the datastore any sort of aggregations query in these cases, as the client is not
+        #   requesting any aggregation data.
+        # - If `SINGLETON_CURSOR` was provide for either `before` or `after`, we also return an empty hash,
+        #   because we know there cannot be any results to return--the cursor is a reference to
+        #   the one and only item in the list, and nothing can exist before or after it.
+        # - Otherwise, we return an aggregatinos hash based on the groupings, computations, and sub-aggregations.
+        def build_agg_hash(filter_interpreter)
+          build_agg_detail(filter_interpreter, field_path: [], parent_queries: [])&.clauses || {}
+        end
+
+        def build_agg_detail(filter_interpreter, field_path:, parent_queries:)
+          return nil if paginator.desired_page_size.zero? || paginator.paginated_from_singleton_cursor?
+          queries = parent_queries + [self] # : ::Array[Query]
+
+          filter_detail(filter_interpreter, field_path) do
+            grouping_adapter.grouping_detail_for(self) do
+              computations_detail.merge(sub_aggregation_detail(filter_interpreter, queries))
+            end
+          end
+        end
+
+        private
+
+        def filter_detail(filter_interpreter, field_path)
+          filtering_field_path = Filtering::FieldPath.of(field_path.filter_map(&:name_in_index))
+          filter_clause = filter_interpreter.build_query([filter].compact, from_field_path: filtering_field_path)
+
+          inner_detail = yield
+
+          return inner_detail if filter_clause.nil?
+          key = "#{name}:filtered"
+
+          clause = {
+            key => {
+              "filter" => filter_clause,
+              "aggs" => inner_detail.clauses
+            }.compact
+          }
+
+          inner_meta = inner_detail.meta
+          meta =
+            if (buckets_path = inner_detail.meta["buckets_path"])
+              # In this case, we have some grouping aggregations applied, and the response will include a `buckets` array.
+              # Here we are prefixing the `buckets_path` with the `key` used for our filter aggregation to maintain its accuracy.
+              inner_meta.merge({"buckets_path" => [key] + buckets_path})
+            else
+              # In this case, no grouping aggregations have been applied, and the response will _not_ have a `buckets` array.
+              # Instead, we'll need to treat the single unbucketed aggregation as a single bucket. To indicate that, we use
+              # `bucket_path` (singular) rather than `buckets_path` (plural).
+              inner_meta.merge({"bucket_path" => [key]})
+            end
+
+          AggregationDetail.new(clause, meta)
+        end
+
+        def computations_detail
+          build_inner_aggregation_detail(computations) do |computation|
+            {computation.key(aggregation_name: name) => computation.clause}
+          end
+        end
+
+        def sub_aggregation_detail(filter_interpreter, parent_queries)
+          build_inner_aggregation_detail(sub_aggregations.values) do |sub_agg|
+            sub_agg.build_agg_hash(filter_interpreter, parent_queries: parent_queries)
+          end
+        end
+
+        def build_inner_aggregation_detail(collection, &block)
+          collection.map(&block).reduce({}, :merge)
+        end
+      end
+
+      # The details of an aggregation level, including the `aggs` clauses themselves and `meta`
+      # that we want echoed back to us in the response for the aggregation level.
+      AggregationDetail = ::Data.define(
+        # Aggregation clauses that would go under `aggs.
+        :clauses,
+        # Custom metadata that will be echoed back to us in the response.
+        # https://www.elastic.co/guide/en/elasticsearch/reference/8.11/search-aggregations.html#add-metadata-to-an-agg
+        :meta
+      ) do
+        # @implements AggregationDetail
+
+        # Wraps this aggregation detail in another aggregation layer for the given `grouping`,
+        # so that we can easily build up the necessary multi-level aggregation structure.
+        def wrap_with_grouping(grouping, query:)
+          agg_key = grouping.key
+          extra_inner_meta = grouping.inner_meta.merge({
+            # The response just includes tuples of values for the key of each bucket. We need to know what fields those
+            # values come from, and this `meta` field  indicates that.
+            "grouping_fields" => [agg_key]
+          })
+
+          inner_agg_hash = {
+            "aggs" => (clauses unless (clauses || {}).empty?),
+            "meta" => meta.merge(extra_inner_meta)
+          }.compact
+
+          missing_bucket_inner_agg_hash = inner_agg_hash.key?("aggs") ? inner_agg_hash : {} # : ::Hash[::String, untyped]
+
+          AggregationDetail.new(
+            {
+              agg_key => grouping.non_composite_clause_for(query).merge(inner_agg_hash),
+
+              # Here we include a `missing` aggregation as a sibling to the main grouping aggregation. We do this
+              # so that we get a bucket of documents that have `null` values for the field we are grouping on, in
+              # order to provide the same behavior as the `CompositeGroupingAdapter` (which uses the built-in
+              # `missing_bucket` option).
+              #
+              # To work correctly, we need to include this `missing` aggregation as a sibling at _every_ level of
+              # the aggregation structure, and the `missing` aggregation needs the same child aggregations as the
+              # main grouping aggregation has. Given the recursive nature of how this is applied, this results in
+              # a fairly complex structure, even though conceptually the idea behind this isn't _too_ bad.
+              Key.missing_value_bucket_key(agg_key) => {
+                "missing" => {"field" => grouping.encoded_index_field_path}
+              }.merge(missing_bucket_inner_agg_hash)
+            },
+            {"buckets_path" => [agg_key]}
+          )
+        end
+      end
+    end
+  end
+end

data/lib/elastic_graph/graphql/aggregation/query_adapter.rb

@@ -0,0 +1,345 @@
+# Copyright 2024 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/graphql/aggregation/composite_grouping_adapter"
+require "elastic_graph/graphql/aggregation/computation"
+require "elastic_graph/graphql/aggregation/date_histogram_grouping"
+require "elastic_graph/graphql/aggregation/field_term_grouping"
+require "elastic_graph/graphql/aggregation/nested_sub_aggregation"
+require "elastic_graph/graphql/aggregation/path_segment"
+require "elastic_graph/graphql/aggregation/query"
+require "elastic_graph/graphql/aggregation/script_term_grouping"
+require "elastic_graph/graphql/schema/arguments"
+require "elastic_graph/support/memoizable_data"
+
+module ElasticGraph
+  class GraphQL
+    module Aggregation
+      # Responsible for taking in the incoming GraphQL request context, arguments, and the GraphQL
+      # schema and directives and populating the `aggregations` portion of `query`.
+      class QueryAdapter < Support::MemoizableData.define(:schema, :config, :filter_args_translator, :runtime_metadata, :sub_aggregation_grouping_adapter)
+        # @dynamic element_names
+        attr_reader :element_names
+
+        def call(query:, lookahead:, args:, field:, context:)
+          return query unless field.type.unwrap_fully.indexed_aggregation?
+
+          aggregations_node = extract_aggregation_node(lookahead, field, context.query)
+          return query unless aggregations_node
+
+          aggregation_query = build_aggregation_query_for(
+            aggregations_node,
+            field: field,
+            grouping_adapter: CompositeGroupingAdapter,
+            # Filters on root aggregations applied to the search query body itself instead of
+            # using a filter aggregation, like sub-aggregations do, so we don't want a filter
+            # aggregation generated here.
+            unfiltered: true
+          )
+
+          query.merge_with(aggregations: {aggregation_query.name => aggregation_query})
+        end
+
+        private
+
+        def after_initialize
+          @element_names = schema.element_names
+        end
+
+        def extract_aggregation_node(lookahead, field, graphql_query)
+          return nil unless (ast_nodes = lookahead.ast_nodes)
+
+          if ast_nodes.size > 1
+            names = ast_nodes.map { |n| "`#{name_of(n)}`" }
+            raise_conflicting_grouping_requirement_selections("`#{lookahead.name}` selection with the same name", names)
+          end
+
+          ::GraphQL::Execution::Lookahead.new(
+            query: graphql_query,
+            ast_nodes: ast_nodes,
+            field: lookahead.field,
+            owner_type: field.parent_type.graphql_type
+          )
+        end
+
+        def build_aggregation_query_for(aggregations_node, field:, grouping_adapter:, nested_path: [], unfiltered: false)
+          aggregation_name = name_of(_ = aggregations_node.ast_nodes.first)
+
+          # Get the AST node for the `nodes` subfield (e.g. from `fooAggregations { nodes { ... } }`)
+          nodes_node = selection_above_grouping_fields(aggregations_node, element_names.nodes, aggregation_name)
+
+          # Also get the AST node for `edges.node` (e.g. from `fooAggregations { edges { node { ... } } }`)
+          edges_node_node = [element_names.edges, element_names.node].reduce(aggregations_node) do |node, sub_selection|
+            selection_above_grouping_fields(node, sub_selection, aggregation_name)
+          end
+
+          # ...and then determine which one is being used for nodes.
+          node_node =
+            if nodes_node.selected? && edges_node_node.selected?
+              raise_conflicting_grouping_requirement_selections("node selection", ["`#{element_names.nodes}`", "`#{element_names.edges}.#{element_names.node}`"])
+            elsif !nodes_node.selected?
+              edges_node_node
+            else
+              nodes_node
+            end
+
+          count_detail_node = node_node.selection(element_names.count_detail)
+          needs_doc_count_error =
+            # We need to know what the error is to determine if the approximate count is in fact the exact count.
+            count_detail_node.selects?(element_names.exact_value) ||
+            # We need to know what the error is to determine the upper bound on the count.
+            count_detail_node.selects?(element_names.upper_bound)
+
+          unless unfiltered
+            filter = filter_args_translator.translate_filter_args(field: field, args: field.args_to_schema_form(aggregations_node.arguments))
+          end
+
+          Query.new(
+            name: aggregation_name,
+            groupings: build_groupings_from(node_node, aggregation_name, from_field_path: nested_path),
+            computations: build_computations_from(node_node, from_field_path: nested_path),
+            sub_aggregations: build_sub_aggregations_from(node_node, parent_nested_path: nested_path),
+            needs_doc_count: count_detail_node.selected? || node_node.selects?(element_names.count),
+            needs_doc_count_error: needs_doc_count_error,
+            paginator: build_paginator_for(aggregations_node),
+            filter: filter,
+            grouping_adapter: grouping_adapter
+          )
+        end
+
+        # Helper method for dealing with lookahead selections above the grouping fields. If the caller selects
+        # such a field multiple times (e.g. with aliases) it leads to conflicting grouping requirements, so we
+        # do not allow it.
+        def selection_above_grouping_fields(node, sub_selection_name, aggregation_name)
+          node.selection(sub_selection_name).tap do |nested_node|
+            ast_nodes = nested_node.ast_nodes || []
+            if ast_nodes.size > 1
+              names = ast_nodes.map { |n| "`#{name_of(n)}`" }
+              raise_conflicting_grouping_requirement_selections("`#{sub_selection_name}` selection under `#{aggregation_name}`", names)
+            end
+          end
+        end
+
+        def build_clauses_from(parent_node, &block)
+          get_children_nodes(parent_node).flat_map do |child_node|
+            transform_node_to_clauses(child_node, &block)
+          end.to_set
+        end
+
+        # Takes a `GraphQL::Execution::Lookahead` node and returns an array of children
+        # lookahead nodes, excluding nodes for introspection fields.
+        def get_children_nodes(node)
+          node.selections.reject do |child|
+            child.field.introspection?
+          end
+        end
+
+        # Takes a `GraphQL::Execution::Lookahead` node that conforms to our aggregate field
+        # conventions (`some_field: {Type}Metric`) and returns a Hash compatible with the `aggregations`
+        # argument to `DatastoreQuery.new`.
+        def transform_node_to_clauses(node, parent_path: [], &clause_builder)
+          field = field_from_node(node)
+          field_path = parent_path + [PathSegment.for(field: field, lookahead: node)]
+
+          clause_builder.call(node, field, field_path) || get_children_nodes(node).flat_map do |embedded_field|
+            transform_node_to_clauses(embedded_field, parent_path: field_path, &clause_builder)
+          end
+        end
+
+        def build_computations_from(node_node, from_field_path: [])
+          aggregated_values_node = node_node.selection(element_names.aggregated_values)
+
+          build_clauses_from(aggregated_values_node) do |node, field, field_path|
+            if field.aggregated?
+              field_path = from_field_path + field_path
+
+              get_children_nodes(node).map do |fn_node|
+                computed_field = field_from_node(fn_node)
+                computation_detail = field_from_node(fn_node).computation_detail # : SchemaArtifacts::RuntimeMetadata::ComputationDetail
+
+                Aggregation::Computation.new(
+                  source_field_path: field_path,
+                  computed_index_field_name: computed_field.name_in_index.to_s,
+                  detail: computation_detail
+                )
+              end
+            end
+          end
+        end
+
+        def build_groupings_from(node_node, aggregation_name, from_field_path: [])
+          grouped_by_node = selection_above_grouping_fields(node_node, element_names.grouped_by, aggregation_name)
+
+          build_clauses_from(grouped_by_node) do |node, field, field_path|
+            field_path = from_field_path + field_path
+
+            # New date/time grouping API (DateGroupedBy, DateTimeGroupedBy)
+            if field.type.elasticgraph_category == :date_grouped_by_object
+              date_time_groupings_from(field_path: field_path, node: node)
+
+            elsif !field.type.object?
+              case field.type.name
+              # Legacy date grouping API
+              when :Date
+                legacy_date_histogram_groupings_from(
+                  field_path: field_path,
+                  node: node,
+                  get_time_zone: ->(args) {},
+                  get_offset: ->(args) { args[element_names.offset_days]&.then { |days| "#{days}d" } }
+                )
+              # Legacy datetime grouping API
+              when :DateTime
+                legacy_date_histogram_groupings_from(
+                  field_path: field_path,
+                  node: node,
+                  get_time_zone: ->(args) { args.fetch(element_names.time_zone) },
+                  get_offset: ->(args) { datetime_offset_from(node, args) }
+                )
+              # Non-date/time grouping
+              else
+                [FieldTermGrouping.new(field_path: field_path)]
+              end
+            end
+          end
+        end
+
+        # Given a `GraphQL::Execution::Lookahead` node, returns the corresponding `Schema::Field`
+        def field_from_node(node)
+          schema.field_named(node.owner_type.graphql_name, node.field.name)
+        end
+
+        # Returns an array of `...Grouping`, one for each child node (`as_date_time`, `as_date`, etc).
+        def date_time_groupings_from(field_path:, node:)
+          get_children_nodes(node).map do |child_node|
+            schema_args = Schema::Arguments.to_schema_form(child_node.arguments, child_node.field)
+            # Because `DateGroupedBy` doesn't have a `timeZone` argument, and we want to reuse the same
+            # script for both `Date` and `DateTime`, we fall back to "UTC" here.
+            time_zone = schema_args[element_names.time_zone] || "UTC"
+            child_field_path = field_path + [PathSegment.for(lookahead: child_node)]
+
+            if child_node.field.name == element_names.as_day_of_week
+              ScriptTermGrouping.new(
+                field_path: child_field_path,
+                script_id: runtime_metadata.static_script_ids_by_scoped_name.fetch("field/as_day_of_week"),
+                params: {
+                  "offset_ms" => datetime_offset_as_ms_from(child_node, schema_args),
+                  "time_zone" => time_zone
+                }
+              )
+            elsif child_node.field.name == element_names.as_time_of_day
+              ScriptTermGrouping.new(
+                field_path: child_field_path,
+                script_id: runtime_metadata.static_script_ids_by_scoped_name.fetch("field/as_time_of_day"),
+                params: {
+                  "interval" => interval_from(child_node, schema_args, interval_unit_key: element_names.truncation_unit),
+                  "offset_ms" => datetime_offset_as_ms_from(child_node, schema_args),
+                  "time_zone" => time_zone
+                }
+              )
+            else
+              DateHistogramGrouping.new(
+                field_path: child_field_path,
+                interval: interval_from(child_node, schema_args, interval_unit_key: element_names.truncation_unit),
+                offset: datetime_offset_from(child_node, schema_args),
+                time_zone: time_zone
+              )
+            end
+          end
+        end
+
+        def legacy_date_histogram_groupings_from(field_path:, node:, get_time_zone:, get_offset:)
+          schema_args = Schema::Arguments.to_schema_form(node.arguments, node.field)
+
+          [DateHistogramGrouping.new(
+            field_path: field_path,
+            interval: interval_from(node, schema_args, interval_unit_key: element_names.granularity),
+            time_zone: get_time_zone.call(schema_args),
+            offset: get_offset.call(schema_args)
+          )]
+        end
+
+        # Figure out the Date histogram grouping interval for the given node based on the `grouped_by` argument.
+        # Until `legacy_grouping_schema` is removed, we need to check both `granularity` and `truncation_unit`.
+        def interval_from(node, schema_args, interval_unit_key:)
+          enum_type_name = node.field.arguments.fetch(interval_unit_key).type.unwrap.graphql_name
+          enum_value_name = schema_args.fetch(interval_unit_key)
+          enum_value = schema.type_named(enum_type_name).enum_value_named(enum_value_name)
+
+          _ = enum_value.runtime_metadata.datastore_value
+        end
+
+        def datetime_offset_from(node, schema_args)
+          if (unit_name = schema_args.dig(element_names.offset, element_names.unit))
+            enum_value = enum_value_from_offset(node, unit_name)
+            amount = schema_args.fetch(element_names.offset).fetch(element_names.amount)
+            "#{amount}#{enum_value.runtime_metadata.datastore_abbreviation}"
+          end
+        end
+
+        # Convert from amount and unit to milliseconds, using runtime metadata `datastore_value`
+        def datetime_offset_as_ms_from(node, schema_args)
+          unit_name = schema_args.dig(element_names.offset, element_names.unit)
+          return 0 unless unit_name
+
+          amount = schema_args.fetch(element_names.offset).fetch(element_names.amount)
+          enum_value = enum_value_from_offset(node, unit_name)
+
+          amount * enum_value.runtime_metadata.datastore_value
+        end
+
+        def enum_value_from_offset(node, unit_name)
+          offset_input_type = node.field.arguments.fetch(element_names.offset).type.unwrap # : ::GraphQL::Schema::InputObject
+          enum_type_name = offset_input_type.arguments.fetch(element_names.unit).type.unwrap.graphql_name
+          schema.type_named(enum_type_name).enum_value_named(unit_name)
+        end
+
+        def name_of(ast_node)
+          ast_node.alias || ast_node.name
+        end
+
+        def build_sub_aggregations_from(node_node, parent_nested_path: [])
+          build_clauses_from(node_node.selection(element_names.sub_aggregations)) do |node, field, field_path|
+            if field.type.elasticgraph_category == :nested_sub_aggregation_connection
+              nested_path = parent_nested_path + field_path
+              nested_sub_agg = NestedSubAggregation.new(
+                nested_path: nested_path,
+                query: build_aggregation_query_for(
+                  node,
+                  field: field,
+                  grouping_adapter: sub_aggregation_grouping_adapter,
+                  nested_path: nested_path
+                )
+              )
+
+              [[nested_sub_agg.query.name, nested_sub_agg]]
+            end
+          end.to_h
+        end
+
+        def build_paginator_for(node)
+          args = field_from_node(node).args_to_schema_form(node.arguments)
+
+          DatastoreQuery::Paginator.new(
+            first: args[element_names.first],
+            after: args[element_names.after],
+            last: args[element_names.last],
+            before: args[element_names.before],
+            default_page_size: config.default_page_size,
+            max_page_size: config.max_page_size,
+            schema_element_names: schema.element_names
+          )
+        end
+
+        def raise_conflicting_grouping_requirement_selections(more_than_one_description, paths)
+          raise ::GraphQL::ExecutionError, "Cannot have more than one #{more_than_one_description} " \
+            "(#{paths.join(", ")}), because that could lead to conflicting grouping requirements."
+        end
+      end
+    end
+  end
+end