elasticgraph-graphql 0.18.0.0

data/lib/elastic_graph/graphql/datastore_query.rb

@@ -0,0 +1,372 @@
+# 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/error"
+require "elastic_graph/graphql/aggregation/query"
+require "elastic_graph/graphql/aggregation/query_optimizer"
+require "elastic_graph/graphql/decoded_cursor"
+require "elastic_graph/graphql/datastore_response/search_response"
+require "elastic_graph/graphql/filtering/filter_interpreter"
+require "elastic_graph/support/memoizable_data"
+
+module ElasticGraph
+  class GraphQL
+    # An immutable class that represents a datastore query. Since this represents
+    # a datastore query, and not a GraphQL query, all the data in it is modeled
+    # in datastore terms, not GraphQL terms. For example, any field names in a
+    # `Query` should be references to index fields, not GraphQL fields.
+    #
+    # Filters are modeled as a `Set` of filtering hashes. While we usually expect only
+    # a single `filter` hash, modeling it as a set makes it easy for us to support
+    # merging queries. The datastore knows how to apply multiple `must` clauses that
+    # apply to the same field, giving us the exact semantics we want in such a situation
+    # with minimal effort.
+    class DatastoreQuery < Support::MemoizableData.define(
+      :total_document_count_needed, :aggregations, :logger, :filter_interpreter, :routing_picker,
+      :index_expression_builder, :default_page_size, :search_index_definitions, :max_page_size,
+      :filters, :sort, :document_pagination, :requested_fields, :individual_docs_needed,
+      :monotonic_clock_deadline, :schema_element_names
+    ) {
+      def initialize(
+        filter: nil,
+        filters: nil,
+        sort: nil,
+        document_pagination: nil,
+        aggregations: nil,
+        requested_fields: nil,
+        individual_docs_needed: false,
+        total_document_count_needed: false,
+        monotonic_clock_deadline: nil,
+        **kwargs
+      )
+        # Deal with `:filter` vs `:filters` input and normalize it to a single `filters` set.
+        filters = ::Set.new(filters || [])
+        filters << filter if filter && !filter.empty?
+        filters.freeze
+
+        aggregations ||= {}
+        requested_fields ||= []
+
+        super(
+          filters: filters,
+          sort: sort || [],
+          document_pagination: document_pagination || {},
+          aggregations: aggregations,
+          requested_fields: requested_fields.to_set,
+          individual_docs_needed: individual_docs_needed || !requested_fields.empty?,
+          total_document_count_needed: total_document_count_needed || aggregations.values.any?(&:needs_total_doc_count?),
+          monotonic_clock_deadline: monotonic_clock_deadline,
+          **kwargs
+        )
+
+        if search_index_definitions.empty?
+          raise SearchFailedError, "Query is invalid, since it contains no `search_index_definitions`."
+        end
+      end
+    }
+      # Load these files after the `Query` class has been defined, to avoid
+      # `TypeError: superclass mismatch for class Query`
+      require "elastic_graph/graphql/datastore_query/document_paginator"
+      require "elastic_graph/graphql/datastore_query/index_expression_builder"
+      require "elastic_graph/graphql/datastore_query/paginator"
+      require "elastic_graph/graphql/datastore_query/routing_picker"
+
+      # Performs a list of queries by building a hash of datastore msearch header/body tuples (keyed
+      # by query), yielding them to the caller, and then post-processing the results. The caller is
+      # responsible for returning a hash of responses by query from its block.
+      #
+      # Note that some of the passed queries may not be yielded to the caller; when we can tell
+      # that a query does not have to be sent to the datastore we avoid yielding it from here.
+      # Therefore, the caller should not assume that all queries passed to this method will be
+      # yielded back.
+      #
+      # The return value is a hash of `DatastoreResponse::SearchResponse` objects by query.
+      #
+      # Note: this method uses `send` to work around ruby visibility rules. We do not want
+      # `#decoded_cursor_factory` to be public, as we only need it here, but we cannot access
+      # it from a class method without using `send`.
+      def self.perform(queries)
+        empty_queries, present_queries = queries.partition(&:empty?)
+
+        responses_by_query = Aggregation::QueryOptimizer.optimize_queries(present_queries) do |optimized_queries|
+          header_body_tuples_by_query = optimized_queries.each_with_object({}) do |query, hash|
+            hash[query] = query.to_datastore_msearch_header_and_body
+          end
+
+          yield(header_body_tuples_by_query)
+        end
+
+        empty_responses = empty_queries.each_with_object({}) do |query, hash|
+          hash[query] = DatastoreResponse::SearchResponse::RAW_EMPTY
+        end
+
+        empty_responses.merge(responses_by_query).each_with_object({}) do |(query, response), hash|
+          hash[query] = DatastoreResponse::SearchResponse.build(response, decoded_cursor_factory: query.send(:decoded_cursor_factory))
+        end.tap do |responses_hash|
+          # Callers expect this `perform` method to provide an invariant: the returned hash MUST contain one entry
+          # for each of the `queries` passed in the args. In practice, violating this invariant primarily causes a
+          # problem when the caller uses the `GraphQL::Dataloader` (which happens for every GraphQL request in production...).
+          # However, our tests do not always run queries end-to-end, so this is an added check we want to do, so that
+          # anytime our logic here fails to include a query in the response in any test, we'll be notified of the
+          # problem.
+          expected_queries = queries.to_set
+          actual_queries = responses_hash.keys.to_set
+
+          if expected_queries != actual_queries
+            missing_queries = expected_queries - actual_queries
+            extra_queries = actual_queries - expected_queries
+
+            raise SearchFailedError, "The `responses_hash` does not have the expected set of queries as keys. " \
+              "This can cause problems for the `GraphQL::Dataloader` and suggests a bug in the logic that should be fixed.\n\n" \
+              "Missing queries (#{missing_queries.size}):\n#{missing_queries.map(&:inspect).join("\n")}.\n\n" \
+              "Extra queries (#{extra_queries.size}): #{extra_queries.map(&:inspect).join("\n")}"
+          end
+        end
+      end
+
+      # Merges the provided query, returning a new combined query object.
+      # Both query objects are left unchanged.
+      def merge(other_query)
+        if search_index_definitions != other_query.search_index_definitions
+          raise ElasticGraph::InvalidMergeError, "`search_index_definitions` conflict while merging between " \
+            "#{search_index_definitions} and #{other_query.search_index_definitions}"
+        end
+
+        with(
+          individual_docs_needed: individual_docs_needed || other_query.individual_docs_needed,
+          total_document_count_needed: total_document_count_needed || other_query.total_document_count_needed,
+          filters: filters + other_query.filters,
+          sort: merge_attribute(other_query, :sort),
+          requested_fields: requested_fields + other_query.requested_fields,
+          document_pagination: merge_attribute(other_query, :document_pagination),
+          monotonic_clock_deadline: [monotonic_clock_deadline, other_query.monotonic_clock_deadline].compact.min,
+          aggregations: aggregations.merge(other_query.aggregations)
+        )
+      end
+
+      # Convenience method for merging when you do not have access to an
+      # `DatastoreQuery::Builder`. Allows you to pass the query options you
+      # would like to merge. As with `#merge`, leaves the original query unchanged
+      # and returns a combined query object.
+      def merge_with(**query_options)
+        merge(with(**query_options))
+      end
+
+      # Pairs the multi-search headers and body into a tuple, as per the format required by the datastore:
+      # https://www.elastic.co/guide/en/elasticsearch/reference/current/search-multi-search.html#search-multi-search-api-desc
+      def to_datastore_msearch_header_and_body
+        @to_datastore_msearch_header_and_body ||= [to_datastore_msearch_header, to_datastore_body]
+      end
+
+      # Returns an index_definition expression string to use for searches. This string can specify
+      # multiple indices, use wildcards, etc. For info about what is supported, see:
+      # https://www.elastic.co/guide/en/elasticsearch/reference/current/multi-index.html
+      def search_index_expression
+        @search_index_expression ||= index_expression_builder.determine_search_index_expression(
+          filters,
+          search_index_definitions,
+          # When we have aggregations, we must require indices to search. When we search no indices, the datastore does not return
+          # the standard aggregations response structure, which causes problems.
+          require_indices: !aggregations_datastore_body.empty?
+        ).to_s
+      end
+
+      # Returns the name of the datastore cluster as a String where this query should be setn.
+      # Unless exactly 1 cluster name is found, this method raises a ConfigError.
+      def cluster_name
+        cluster_name = search_index_definitions.map(&:cluster_to_query).uniq
+        return cluster_name.first if cluster_name.size == 1
+        raise ConfigError, "Found different datastore clusters (#{cluster_name}) to query " \
+          "for query targeting indices: #{search_index_definitions}"
+      end
+
+      # Returns a list of unique field paths that should be used for shard routing during searches.
+      #
+      # If a search is filtering on one of these fields, we can optimize the search by routing
+      # it to only the shards containing documents for that routing value.
+      #
+      # Note that this returns a list due to our support for type unions. A unioned type
+      # can be composed of subtypes that have use different shard routing; this will return
+      # the set union of them all.
+      def route_with_field_paths
+        search_index_definitions.map(&:route_with).uniq
+      end
+
+      # The shard routing values used for this search. Can be `nil` if the query will hit all shards.
+      # `[]` means that we are routing to no shards.
+      def shard_routing_values
+        return @shard_routing_values if defined?(@shard_routing_values)
+        routing_values = routing_picker.extract_eligible_routing_values(filters, route_with_field_paths)
+
+        @shard_routing_values ||=
+          if routing_values&.empty? && !aggregations_datastore_body.empty?
+            # If we return an empty array of routing values, no shards will get searched, which causes a problem for aggregations.
+            # When a query includes aggregations, there are normally aggregation structures on the respopnse (even when there are no
+            # search hits to aggregate over!) but if there are no routing values, those aggregation structures will be missing from
+            # the response. It's complex to handle that in our downstream response handling code, so we prefer to force a "fallback"
+            # routing value here to ensure that at least one shard gets searched. Which shard gets searched doesn't matter; the search
+            # filter that led to an empty set of routing values will match on documents on any shard.
+            ["fallback_shard_routing_value"]
+          elsif contains_ignored_values_for_routing?(routing_values)
+            nil
+          else
+            routing_values&.sort # order doesn't matter, but sorting it makes it easier to assert on in our tests.
+          end
+      end
+
+      # Indicates if the query does not need any results from the datastore. As an optimization,
+      # we can reply with a default "empty" response for an empty query.
+      def empty?
+        # If we are searching no indices or routing to an empty set of shards, there is no need to query the datastore at all.
+        # This only happens when our filter processing has deduced that the query will match no results.
+        return true if search_index_expression.empty? || shard_routing_values&.empty?
+
+        datastore_body = to_datastore_body
+        datastore_body.fetch(:size) == 0 && !datastore_body.fetch(:track_total_hits) && aggregations_datastore_body.empty?
+      end
+
+      def inspect
+        description = to_datastore_msearch_header.merge(to_datastore_body).map do |key, value|
+          "#{key}=#{(key == :query) ? "<REDACTED>" : value.inspect}"
+        end.join(" ")
+
+        "#<#{self.class.name} #{description}>"
+      end
+
+      def to_datastore_msearch_header
+        @to_datastore_msearch_header ||= {index: search_index_expression, routing: shard_routing_values&.join(",")}.compact
+      end
+
+      # `DatastoreQuery` objects are used as keys in a hash. Computing `#hash` can be expensive (given how many fields
+      # an `DatastoreQuery` has) and it's safe to cache since `DatastoreQuery` instances are immutable, so we memoize it
+      # here. We've observed this making a very noticeable difference in our test suite runtime.
+      def hash
+        @hash ||= super
+      end
+
+      def document_paginator
+        @document_paginator ||= DocumentPaginator.new(
+          sort_clauses: sort_with_tiebreaker,
+          individual_docs_needed: individual_docs_needed,
+          total_document_count_needed: total_document_count_needed,
+          decoded_cursor_factory: decoded_cursor_factory,
+          schema_element_names: schema_element_names,
+          paginator: Paginator.new(
+            default_page_size: default_page_size,
+            max_page_size: max_page_size,
+            first: document_pagination[:first],
+            after: document_pagination[:after],
+            last: document_pagination[:last],
+            before: document_pagination[:before],
+            schema_element_names: schema_element_names
+          )
+        )
+      end
+
+      private
+
+      def merge_attribute(other_query, attribute)
+        value = public_send(attribute)
+        other_value = other_query.public_send(attribute)
+
+        if value.empty?
+          other_value
+        elsif other_value.empty?
+          value
+        elsif value == other_value
+          value
+        else
+          logger.warn("Tried to merge two queries that both define `#{attribute}`, using the value from the query being merged: #{value}, #{other_value}")
+          other_value
+        end
+      end
+
+      TIEBREAKER_SORT_CLAUSES = [{"id" => {"order" => "asc"}}].freeze
+
+      # We want to use `id` as a tiebreaker ONLY when `id` isn't explicitly specified as a sort field
+      def sort_with_tiebreaker
+        @sort_with_tiebreaker ||= remove_duplicate_sort_clauses(sort + TIEBREAKER_SORT_CLAUSES)
+      end
+
+      def remove_duplicate_sort_clauses(sort_clauses)
+        seen_fields = Set.new
+        sort_clauses.select do |clause|
+          clause.keys.all? { |key| seen_fields.add?(key) }
+        end
+      end
+
+      def decoded_cursor_factory
+        @decoded_cursor_factory ||= DecodedCursor::Factory.from_sort_list(sort_with_tiebreaker)
+      end
+
+      def contains_ignored_values_for_routing?(routing_values)
+        ignored_values_for_routing.intersect?(routing_values.to_set) if routing_values
+      end
+
+      def ignored_values_for_routing
+        @ignored_values_for_routing ||= search_index_definitions.flat_map { |i| i.ignored_values_for_routing.to_a }.to_set
+      end
+
+      def to_datastore_body
+        @to_datastore_body ||= aggregations_datastore_body
+          .merge(document_paginator.to_datastore_body)
+          .merge({query: filter_interpreter.build_query(filters)}.compact)
+          .merge({_source: source})
+      end
+
+      def aggregations_datastore_body
+        @aggregations_datastore_body ||= begin
+          aggs = aggregations
+            .values
+            .map { |agg| agg.build_agg_hash(filter_interpreter) }
+            .reduce({}, :merge)
+
+          aggs.empty? ? {} : {aggs: aggs}
+        end
+      end
+
+      # Make our query as efficient as possible by limiting what parts of `_source` we fetch.
+      # For an id-only query (or a query that has no requested fields) we don't need to fetch `_source`
+      # at all--which means the datastore can avoid decompressing the _source field. Otherwise,
+      # we only ask for the fields we need to return.
+      def source
+        requested_source_fields = requested_fields - ["id"]
+        return false if requested_source_fields.empty?
+        # Merging in requested_fields as _source:{includes:} based on Elasticsearch documentation:
+        # https://www.elastic.co/guide/en/elasticsearch/reference/current/mapping-source-field.html#include-exclude
+        {includes: requested_source_fields.to_a}
+      end
+
+      # Encapsulates dependencies of `Query`, giving us something we can expose off of `application`
+      # to build queries when desired.
+      class Builder < Support::MemoizableData.define(:runtime_metadata, :logger, :query_defaults)
+        def self.with(runtime_metadata:, logger:, **query_defaults)
+          new(runtime_metadata: runtime_metadata, logger: logger, query_defaults: query_defaults)
+        end
+
+        def routing_picker
+          @routing_picker ||= RoutingPicker.new(schema_names: runtime_metadata.schema_element_names)
+        end
+
+        def index_expression_builder
+          @index_expression_builder ||= IndexExpressionBuilder.new(schema_names: runtime_metadata.schema_element_names)
+        end
+
+        def new_query(**options)
+          DatastoreQuery.new(
+            routing_picker: routing_picker,
+            index_expression_builder: index_expression_builder,
+            logger: logger,
+            schema_element_names: runtime_metadata.schema_element_names,
+            **query_defaults.merge(options)
+          )
+        end
+      end
+    end
+  end
+end

data/lib/elastic_graph/graphql/datastore_response/document.rb

@@ -0,0 +1,78 @@
+# 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/decoded_cursor"
+require "elastic_graph/support/memoizable_data"
+require "forwardable"
+
+module ElasticGraph
+  class GraphQL
+    module DatastoreResponse
+      # Represents a document fetched from the datastore. Exposes both the raw metadata
+      # provided by the datastore and the doc payload itself. In addition, you can treat
+      # it just like a document hash using `#[]` or `#fetch`.
+      Document = Support::MemoizableData.define(:raw_data, :payload, :decoded_cursor_factory) do
+        # @implements Document
+        extend Forwardable
+        def_delegators :payload, :[], :fetch
+
+        def self.build(raw_data, decoded_cursor_factory: DecodedCursor::Factory::Null)
+          source = raw_data.fetch("_source") do
+            {} # : ::Hash[::String, untyped]
+          end
+
+          new(
+            raw_data: raw_data,
+            # Since we no longer fetch _source for id only queries, merge id into _source to take care of that case
+            payload: source.merge("id" => raw_data["_id"]),
+            decoded_cursor_factory: decoded_cursor_factory
+          )
+        end
+
+        def self.with_payload(payload)
+          build({"_source" => payload})
+        end
+
+        def index_name
+          raw_data["_index"]
+        end
+
+        def index_definition_name
+          index_name.split(ROLLOVER_INDEX_INFIX_MARKER).first # : ::String
+        end
+
+        def id
+          raw_data["_id"]
+        end
+
+        def sort
+          raw_data["sort"]
+        end
+
+        def version
+          payload["version"]
+        end
+
+        def cursor
+          @cursor ||= decoded_cursor_factory.build(raw_data.fetch("sort"))
+        end
+
+        def datastore_path
+          # Path based on this API:
+          # https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-get.html
+          "/#{index_name}/_doc/#{id}".squeeze("/")
+        end
+
+        def to_s
+          "#<#{self.class.name} #{datastore_path}>"
+        end
+        alias_method :inspect, :to_s
+      end
+    end
+  end
+end

data/lib/elastic_graph/graphql/datastore_response/search_response.rb

@@ -0,0 +1,79 @@
+# 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/error"
+require "elastic_graph/graphql/decoded_cursor"
+require "elastic_graph/graphql/datastore_response/document"
+require "forwardable"
+
+module ElasticGraph
+  class GraphQL
+    module DatastoreResponse
+      # Represents a search response from the datastore. Exposes both the raw metadata
+      # provided by the datastore and the collection of documents. Can be treated as a
+      # collection of documents when you don't care about the metadata.
+      class SearchResponse < ::Data.define(:raw_data, :metadata, :documents, :total_document_count)
+        include Enumerable
+        extend Forwardable
+
+        def_delegators :documents, :each, :to_a, :size, :empty?
+
+        EXCLUDED_METADATA_KEYS = %w[hits aggregations].freeze
+
+        def self.build(raw_data, decoded_cursor_factory: DecodedCursor::Factory::Null)
+          documents = raw_data.fetch("hits").fetch("hits").map do |doc|
+            Document.build(doc, decoded_cursor_factory: decoded_cursor_factory)
+          end
+
+          metadata = raw_data.except(*EXCLUDED_METADATA_KEYS)
+          metadata["hits"] = raw_data.fetch("hits").except("hits")
+
+          # `hits.total` is exposed as an object like:
+          #
+          # {
+          #   "value" => 200,
+          #   "relation" => "eq", # or "gte"
+          # }
+          #
+          # This allows it to provide a lower bound on the number of hits, rather than having
+          # to give an exact count. We may want to handle the `gte` case differently at some
+          # point but for now we just use the value as-is.
+          #
+          # In the case where `track_total_hits` flag is set to `false`, `hits.total` field will be completely absent.
+          # This means the client intentionally chose not to query the total doc count, and `total_document_count` will be nil.
+          # In this case, we will throw an exception if the client later tries to access `total_document_count`.
+          total_document_count = metadata.dig("hits", "total", "value")
+
+          new(
+            raw_data: raw_data,
+            metadata: metadata,
+            documents: documents,
+            total_document_count: total_document_count
+          )
+        end
+
+        # Benign empty response that can be used in place of datastore response errors as needed.
+        RAW_EMPTY = {"hits" => {"hits" => [], "total" => {"value" => 0}}}.freeze
+        EMPTY = build(RAW_EMPTY)
+
+        def docs_description
+          (documents.size < 3) ? documents.inspect : "[#{documents.first}, ..., #{documents.last}]"
+        end
+
+        def total_document_count
+          super || raise(CountUnavailableError, "#{__method__} is unavailable; set `query.total_document_count_needed = true` to make it available")
+        end
+
+        def to_s
+          "#<#{self.class.name} size=#{documents.size} #{docs_description}>"
+        end
+        alias_method :inspect, :to_s
+      end
+    end
+  end
+end

data/lib/elastic_graph/graphql/datastore_search_router.rb

@@ -0,0 +1,151 @@
+# 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/constants"
+require "elastic_graph/error"
+require "elastic_graph/graphql/datastore_response/search_response"
+require "elastic_graph/graphql/query_details_tracker"
+require "elastic_graph/support/threading"
+
+module ElasticGraph
+  class GraphQL
+    # Responsible for routing datastore search requests to the appropriate cluster and index.
+    class DatastoreSearchRouter
+      def initialize(
+        datastore_clients_by_name:,
+        logger:,
+        monotonic_clock:,
+        config:
+      )
+        @datastore_clients_by_name = datastore_clients_by_name
+        @logger = logger
+        @monotonic_clock = monotonic_clock
+        @config = config
+      end
+
+      # Sends the datastore a multi-search request based on the given queries.
+      # Returns a hash of responses keyed by the query.
+      def msearch(queries, query_tracker: QueryDetailsTracker.empty)
+        DatastoreQuery.perform(queries) do |header_body_tuples_by_query|
+          # Here we set a client-side timeout, which causes the client to give up and close the connection.
+          # According to [1]--"We have a new way to cancel search requests efficiently from the client
+          # in 7.4 (by closing the underlying http channel)"--this should cause the server to stop
+          # executing the search, and more importantly, gives us a strictly enforced timeout.
+          #
+          # In addition, the datastore supports a `timeout` option on a search body, but this timeout is
+          # "best effort", applies to each shard (and not to the overall search request), and only interrupts
+          # certain kinds of operations. [2] and [3] below have more info.
+          #
+          # Note that I have not been able to observe this `timeout` on a search body ever working
+          # as documented. In our test suite, none of the slow queries I have tried (both via
+          # slow aggregation query and a slow script) have ever aborted early when that option is
+          # set. In Kibana in production, @bsorbo observed it aborting a `search` request early
+          # (but not necessarily an `msearch` request...), but even then, the response said `timed_out: false`!
+          # Other people ([4]) have reported observing timeout having no effect on msearch requests.
+          #
+          # So, the client-side timeout is the main one we want here, and for now we are not using the
+          # datastore search `timeout` option at all.
+          #
+          # For more info, see:
+          #
+          # [1] https://github.com/elastic/elasticsearch/issues/47716
+          # [2] https://github.com/elastic/elasticsearch/pull/51858
+          # [3] https://www.elastic.co/guide/en/elasticsearch/guide/current/_search_options.html#_timeout_2
+          # [4] https://discuss.elastic.co/t/timeouts-ignored-in-multisearch/23673
+
+          # Unfortunately, the Elasticsearch/OpenSearch clients don't support setting a per-request client-side timeout,
+          # even though Faraday (the underlying HTTP client) does. To work around this, we pass our desired
+          # timeout in a specific header that the `SupportTimeouts` Faraday middleware will use.
+          headers = {TIMEOUT_MS_HEADER => msearch_request_timeout_from(queries)}.compact
+
+          queries_and_header_body_tuples_by_datastore_client = header_body_tuples_by_query.group_by do |(query, header_body_tuples)|
+            @datastore_clients_by_name.fetch(query.cluster_name)
+          end
+
+          datastore_query_started_at = @monotonic_clock.now_in_ms
+
+          server_took_and_results = Support::Threading.parallel_map(queries_and_header_body_tuples_by_datastore_client) do |datastore_client, query_and_header_body_tuples_for_cluster|
+            queries_for_cluster, header_body_tuples = query_and_header_body_tuples_for_cluster.transpose
+            msearch_body = header_body_tuples.flatten(1)
+            response = datastore_client.msearch(body: msearch_body, headers: headers)
+            debug_query(query: msearch_body, response: response)
+            ordered_responses = response.fetch("responses")
+            [response["took"], queries_for_cluster.zip(ordered_responses)]
+          end
+
+          query_tracker.record_datastore_query_duration_ms(
+            client: @monotonic_clock.now_in_ms - datastore_query_started_at,
+            server: server_took_and_results.map(&:first).compact.max
+          )
+
+          server_took_and_results.flat_map(&:last).to_h.tap do |responses_by_query|
+            log_shard_failure_if_necessary(responses_by_query)
+            raise_search_failed_if_any_failures(responses_by_query)
+          end
+        end
+      end
+
+      private
+
+      # Prefix tests with `DEBUG_QUERY=1 ...` or run `export DEBUG_QUERY=1` to print the actual
+      # Elasticsearch/OpenSearch query and response. This is particularly useful for adding new specs.
+      def debug_query(**debug_messages)
+        return unless ::ENV["DEBUG_QUERY"]
+
+        formatted_messages = debug_messages.map do |key, msg|
+          "#{key.to_s.upcase}:\n#{::JSON.pretty_generate(msg)}\n"
+        end.join("\n")
+        puts "\n#{formatted_messages}\n\n"
+      end
+
+      def msearch_request_timeout_from(queries)
+        return nil unless (min_query_deadline = queries.map(&:monotonic_clock_deadline).compact.min)
+
+        (min_query_deadline - @monotonic_clock.now_in_ms).tap do |timeout|
+          if timeout <= 0
+            raise RequestExceededDeadlineError, "It is already #{timeout.abs} ms past the search deadline."
+          end
+        end
+      end
+
+      def raise_search_failed_if_any_failures(responses_by_query)
+        failures = responses_by_query.each_with_index.select { |(_query, response), _index| response["error"] }
+        return if failures.empty?
+
+        formatted_failures = failures.map do |(query, response), index|
+          # Note: we intentionally omit the body of the request here, because it could contain PII
+          # or other sensitive values that we don't want logged.
+          <<~ERROR
+            #{index + 1}) Header: #{::JSON.generate(query.to_datastore_msearch_header)}
+            #{response.fetch("error").inspect}"
+            On cluster: #{query.cluster_name}
+          ERROR
+        end.join("\n\n")
+
+        raise SearchFailedError, "Got #{failures.size} search failure(s):\n\n#{formatted_failures}"
+      end
+
+      # Examine successful query responses and log any shard failure they encounter
+      def log_shard_failure_if_necessary(responses_by_query)
+        shard_failures = responses_by_query.each_with_index.select do |(query, response), query_numeric_index|
+          (200..299).cover?(response["status"]) && response["_shards"]["failed"] != 0
+        end
+
+        unless shard_failures.empty?
+          formatted_failures = shard_failures.map do |(query, response), query_numeric_index|
+            "Query #{query_numeric_index + 1} against index `#{query.search_index_expression}` on cluster `#{query.cluster_name}`}: " +
+              JSON.pretty_generate(response["_shards"])
+          end.join("\n\n")
+
+          formatted_shard_failures = "The following queries have failed shards: \n\n#{formatted_failures}"
+          @logger.warn(formatted_shard_failures)
+        end
+      end
+    end
+  end
+end