Add automatic X-Opaque-Id headers to datastore msearch requests (#1135)

Author: jwils

elasticgraph-graphql/lib/elastic_graph/graphql/client.rb

@@ -13,12 +13,13 @@ class GraphQL
     # meant to be a friendly/human readable string (such as a service name)
     # where as `source_description` is meant to be an opaque string describing
     # where `name` came from.
-    class Client < Data.define(:source_description, :name)
+    class Client < Data.define(:source_description, :name, :extra_opaque_id_parts)
       # `Data.define` provides the following methods:
-      # @dynamic initialize, name, source_description, with
+      # @dynamic name, source_description, extra_opaque_id_parts, with
 
-      ANONYMOUS = new("(anonymous)", "(anonymous)")
-      ELASTICGRAPH_INTERNAL = new("(ElasticGraphInternal)", "(ElasticGraphInternal)")
+      def initialize(source_description:, name:, extra_opaque_id_parts: [])
+        super
+      end
 
       def description
         if source_description == name
@@ -38,6 +39,9 @@ def resolve(http_request)
           Client::ANONYMOUS
         end
       end
+
+      ANONYMOUS = new("(anonymous)", "(anonymous)")
+      ELASTICGRAPH_INTERNAL = new("(ElasticGraphInternal)", "(ElasticGraphInternal)")
     end
   end
 end

elasticgraph-graphql/lib/elastic_graph/graphql/datastore_search_router.rb

@@ -10,6 +10,7 @@
 require "elastic_graph/errors"
 require "elastic_graph/graphql/datastore_response/search_response"
 require "elastic_graph/graphql/query_details_tracker"
+require "elastic_graph/support/opaque_id"
 require "elastic_graph/support/threading"
 
 module ElasticGraph
@@ -32,7 +33,7 @@ def initialize(
 
       # 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)
+      def msearch(queries, query_tracker: QueryDetailsTracker.empty, opaque_id_parts: ["elasticgraph-graphql"])
         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
@@ -64,6 +65,9 @@ def msearch(queries, query_tracker: QueryDetailsTracker.empty)
           # 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)&.to_s}.compact
+          if (opaque_id = Support::OpaqueID.build_header(opaque_id_parts))
+            headers[OPAQUE_ID_HEADER] = opaque_id
+          end
 
           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)

elasticgraph-graphql/lib/elastic_graph/graphql/query_executor.rb

@@ -59,7 +59,8 @@ def execute(
           client: client,
           context: context.merge({
             monotonic_clock_deadline: timeout_in_ms&.+(start_time_in_ms),
-            elastic_graph_query_tracker: query_tracker
+            elastic_graph_query_tracker: query_tracker,
+            elastic_graph_client: client
           }.compact)
         )
 
@@ -89,7 +90,7 @@ def execute(
           @logger.info({
             "message_type" => "ElasticGraphQueryExecutorQueryDuration",
             "client" => client.name,
-            "query_fingerprint" => fingerprint_for(query),
+            "query_fingerprint" => query.fingerprint,
             "query_name" => query.selected_operation_name,
             "duration_ms" => duration,
             # How long the datastore queries took according to what the datastore itself reported.
@@ -142,7 +143,7 @@ def build_and_execute_query(query_string:, variables:, operation_name:, context:
       def execute_query(query, client:)
         # Log the query before starting to execute it, in case there's a lambda timeout, in which case
         # we won't get any other logged messages for the query.
-        @logger.info "Starting to execute query #{fingerprint_for(query)} for client #{client.description}."
+        @logger.info "Starting to execute query #{query.fingerprint} for client #{client.description}."
 
         query.result
       rescue => ex
@@ -163,11 +164,7 @@ def execute_query(query, client:)
       # the fingerprint to make sure that we at least have some identification information
       # about the query.
       def full_description_of(query)
-        "#{fingerprint_for(query)} #{query.sanitized_query_string}"
-      end
-
-      def fingerprint_for(query)
-        query.query_string ? query.fingerprint : "(no query string)"
+        "#{query.fingerprint} #{query.sanitized_query_string}"
       end
 
       def slo_result_for(query, duration)

elasticgraph-graphql/lib/elastic_graph/graphql/resolvers/query_source.rb

@@ -20,28 +20,48 @@ module Resolvers
       # Note: `NestedRelationshipsSource` implements further optimizations on top of this, and should
       # be used rather than this class when applicable.
       class QuerySource < ::GraphQL::Dataloader::Source
-        def initialize(datastore_router, query_tracker)
+        def initialize(datastore_router, query_tracker, opaque_id_parts = [])
           @datastore_router = datastore_router
           @query_tracker = query_tracker
+          @opaque_id_parts = opaque_id_parts
         end
 
         def fetch(queries)
-          responses_by_query = @datastore_router.msearch(queries, query_tracker: @query_tracker)
+          responses_by_query = @datastore_router.msearch(
+            queries,
+            query_tracker: @query_tracker,
+            opaque_id_parts: @opaque_id_parts
+          )
           queries.map { |q| responses_by_query.fetch(q) }
         end
 
         def self.execute_many(queries, for_context:)
           datastore_router = for_context.fetch(:datastore_search_router)
           query_tracker = for_context.fetch(:elastic_graph_query_tracker)
+          opaque_id_parts = datastore_opaque_id_parts_for(for_context)
           dataloader = for_context.dataloader
 
-          responses = dataloader.with(self, datastore_router, query_tracker).load_all(queries)
+          responses = dataloader.with(self, datastore_router, query_tracker, opaque_id_parts).load_all(queries)
           queries.zip(responses).to_h
         end
 
         def self.execute_one(query, for_context:)
           execute_many([query], for_context: for_context).fetch(query)
         end
+
+        # `QueryExecutor` adds `:elastic_graph_client` to the GraphQL context before
+        # resolver execution begins, so resolver-side datastore queries can reuse the
+        # same client identity in their datastore `X-Opaque-Id` headers.
+        private_class_method def self.datastore_opaque_id_parts_for(for_context)
+          client = for_context.fetch(:elastic_graph_client)
+          graphql_query = for_context.query
+          [
+            "elasticgraph-graphql",
+            "client=#{client.name}",
+            *client.extra_opaque_id_parts,
+            "query=#{graphql_query.fingerprint}"
+          ]
+        end
       end
     end
   end

elasticgraph-graphql/sig/elastic_graph/graphql/client.rbs

@@ -3,12 +3,13 @@ module ElasticGraph
     class Client
       attr_reader name: ::String
       attr_reader source_description: ::String
-      def initialize: (name: ::String, source_description: ::String) -> void
-      def with: (?name: ::String, ?source_description: ::String) -> Client
+      attr_reader extra_opaque_id_parts: ::Array[::String]
+      def initialize: (source_description: ::String, name: ::String, ?extra_opaque_id_parts: ::Array[::String]) -> void
+      def with: (?source_description: ::String, ?name: ::String, ?extra_opaque_id_parts: ::Array[::String]) -> Client
 
       def self.new:
-        (name: ::String, source_description: ::String) -> instance
-      | (::String, ::String) -> instance
+        (source_description: ::String, name: ::String, ?extra_opaque_id_parts: ::Array[::String]) -> instance
+      | (::String, ::String, ?::Array[::String]) -> instance
 
       ANONYMOUS: Client
       ELASTICGRAPH_INTERNAL: Client

elasticgraph-graphql/sig/elastic_graph/graphql/datastore_search_router.rbs

@@ -8,7 +8,11 @@ module ElasticGraph
         config: Config
       ) -> void
 
-      def msearch: (::Array[DatastoreQuery], ?query_tracker: QueryDetailsTracker) -> ::Hash[DatastoreQuery, DatastoreResponse::SearchResponse]
+      def msearch: (
+        ::Array[DatastoreQuery],
+        ?query_tracker: QueryDetailsTracker,
+        ?opaque_id_parts: ::Array[::String?]
+      ) -> ::Hash[DatastoreQuery, DatastoreResponse::SearchResponse]
     end
   end
 end

elasticgraph-graphql/sig/elastic_graph/graphql/query_executor.rbs

@@ -37,7 +37,6 @@ module ElasticGraph
 
       def execute_query: (::GraphQL::Query, client: Client) -> ::GraphQL::Query::Result
       def full_description_of: (::GraphQL::Query) -> ::String
-      def fingerprint_for: (::GraphQL::Query) -> ::String
       def directives_from_query_operation: (::GraphQL::Query) -> ::Hash[::String, ::Hash[::String, untyped]]
       def slo_result_for: (::GraphQL::Query, ::Integer) -> ::String?
     end

elasticgraph-graphql/sig/elastic_graph/graphql/resolvers/query_source.rbs

@@ -2,9 +2,10 @@ module ElasticGraph
   class GraphQL
     module Resolvers
       class QuerySource < ::GraphQL::Dataloader::Source
-        def initialize: (DatastoreSearchRouter, QueryDetailsTracker) -> void
+        def initialize: (DatastoreSearchRouter, QueryDetailsTracker, ?::Array[::String]) -> void
         @datastore_router: DatastoreSearchRouter
         @query_tracker: QueryDetailsTracker
+        @opaque_id_parts: ::Array[::String]
         def fetch: (::Array[DatastoreQuery]) -> ::Array[DatastoreResponse::SearchResponse]
 
         def self.execute_many: (
@@ -16,6 +17,8 @@ module ElasticGraph
           DatastoreQuery,
           for_context: ::GraphQL::Query::Context
         ) -> DatastoreResponse::SearchResponse
+
+        def self.datastore_opaque_id_parts_for: (::GraphQL::Query::Context) -> ::Array[::String]
       end
     end
   end

elasticgraph-graphql/spec/integration/elastic_graph/graphql/resolvers/nested_relationships_source_spec.rb

@@ -7,6 +7,7 @@
 # frozen_string_literal: true
 
 require "elastic_graph/graphql/query_details_tracker"
+require "elastic_graph/graphql/client"
 require "elastic_graph/graphql/resolvers/nested_relationships_source"
 require "graphql"
 require "support/aggregations_helpers"
@@ -217,13 +218,14 @@ def resolve_field(field, *value_sets, **query_args)
 
           ::GraphQL::Dataloader.with_dataloading do |dataloader|
             context = ::GraphQL::Query::Context.new(
-              query: instance_double(::GraphQL::Query),
+              query: instance_double(::GraphQL::Query, fingerprint: "NestedRelationshipsSource/test"),
               schema: graphql.schema.graphql_schema,
               values: {
                 elastic_graph_schema: graphql.schema,
                 dataloader: dataloader,
                 datastore_search_router: graphql.datastore_search_router,
-                elastic_graph_query_tracker: QueryDetailsTracker.empty
+                elastic_graph_query_tracker: QueryDetailsTracker.empty,
+                elastic_graph_client: Client::ANONYMOUS
               }
             )
 

elasticgraph-graphql/spec/support/resolver.rb

@@ -7,6 +7,7 @@
 # frozen_string_literal: true
 
 require "elastic_graph/graphql/query_details_tracker"
+require "elastic_graph/graphql/client"
 require "elastic_graph/graphql/resolvers/query_adapter"
 require "elastic_graph/graphql/resolvers/query_source"
 require "graphql"
@@ -21,13 +22,14 @@ def resolve(type_name, field_name, document = nil, **options)
 
     ::GraphQL::Dataloader.with_dataloading do |dataloader|
       context = ::GraphQL::Query::Context.new(
-        query: nil,
+        query: instance_double(::GraphQL::Query, fingerprint: "ResolverHelperQuery/test"),
         schema: graphql.schema.graphql_schema,
         values: {
           elastic_graph_schema: graphql.schema,
           dataloader: dataloader,
           elastic_graph_query_tracker: query_details_tracker,
-          datastore_search_router: graphql.datastore_search_router
+          datastore_search_router: graphql.datastore_search_router,
+          elastic_graph_client: ElasticGraph::GraphQL::Client::ANONYMOUS
         }
       )