Introduce QueryContext to replace magic GraphQL context keys, closes #1339

Replaces untyped context hash keys (:elastic_graph_schema,
:datastore_search_router, :datastore_query_cache, :elastic_graph_query_tracker,
:monotonic_clock_deadline, :elastic_graph_client, :http_request) with typed
accessors on a QueryContext < GraphQL::Query::Context.

- QueryContext.new_class closes an anonymous subclass over elastic_graph_schema
  and datastore_search_router, registered as the schema's context_class.
- monotonic_clock_deadline/elastic_graph_client/http_request are read-only
  attrs, set once via QueryContext#register_elastic_graph_values (defaults:
  elastic_graph_client: Client::ANONYMOUS, others nil) called by each
  query-building call site right after building the query.
- :visibility_profile stays in the context hash (needed at schema-boot time,
  before a QueryContext exists).
- No dedicated unit spec: existing seam-level specs already give 100%
  coverage, and a direct unit test would pin internal APIs unnecessarily.

Author: myronmarston

elasticgraph-apollo/apollo_tests_implementation/lib/product_resolver.rb

@@ -23,7 +23,7 @@ def initialize(elasticgraph_graphql:, config:)
   def resolve(field:, object:, args:, context:)
     query = @datastore_query_builder.new_query(
       initial_search_index_definitions: [@product_index_def],
-      monotonic_clock_deadline: context[:monotonic_clock_deadline],
+      monotonic_clock_deadline: context.monotonic_clock_deadline,
       client_filters: [{"id" => {"equalToAnyOf" => [args.fetch("id")]}}],
       individual_docs_needed: true,
       request_all_fields: true

elasticgraph-apollo/lib/elastic_graph/apollo/graphql/entities_field_resolver.rb

@@ -25,7 +25,7 @@ def initialize(elasticgraph_graphql:, config:)
         end
 
         def resolve(field:, object:, args:, context:, lookahead:)
-          schema = context.fetch(:elastic_graph_schema)
+          schema = context.elastic_graph_schema
 
           representations = args.fetch("representations").map.with_index do |rep, index|
             try_parse_representation(rep, schema) do |error_description|
@@ -40,7 +40,7 @@ def resolve(field:, object:, args:, context:, lookahead:)
           query_attributes = ElasticGraph::GraphQL::QueryAdapter::RequestedFields
             .new(schema)
             .query_attributes_for(field: field, lookahead: lookahead)
-            .merge(monotonic_clock_deadline: context[:monotonic_clock_deadline])
+            .merge(monotonic_clock_deadline: context.monotonic_clock_deadline)
 
           # Build a separate query per adapter instance since each adapter instance is capable of building
           # a single query that handles all representations assigned to it.

elasticgraph-apollo/lib/elastic_graph/apollo/graphql/service_field_resolver.rb

@@ -18,7 +18,7 @@ def initialize(elasticgraph_graphql:, config:)
         end
 
         def resolve(field:, object:, args:, context:)
-          {"sdl" => service_sdl(context.fetch(:elastic_graph_schema))}
+          {"sdl" => service_sdl(context.elastic_graph_schema)}
         end
 
         private

elasticgraph-apollo/sig/elastic_graph/apollo/graphql/entities_field_resolver.rbs

@@ -43,7 +43,7 @@ module ElasticGraph
           def identify_matching_hit: (
             I,
             R,
-            context: ::GraphQL::Query::Context,
+            context: ElasticGraph::GraphQL::QueryContext,
             index: ::Integer
           ) -> (ElasticGraph::GraphQL::DatastoreResponse::Document | ::Hash[::String, untyped])?
 

elasticgraph-graphql/lib/elastic_graph/graphql/aggregation/query_adapter.rb

@@ -43,7 +43,7 @@ def initialize(config:, filter_args_translator:, runtime_metadata:, sub_aggregat
           def call(query:, lookahead:, args:, field:, context:)
             return query unless field.type.unwrap_fully.indexed_aggregation?
 
-            @build_adapter.call(context.fetch(:elastic_graph_schema)).call(
+            @build_adapter.call(context.elastic_graph_schema).call(
               query: query,
               lookahead: lookahead,
               args: args,

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

@@ -57,6 +57,7 @@ def process(request, max_timeout_in_ms: nil, start_time_in_ms: @monotonic_clock.
             client: client_or_response,
             timeout_in_ms: parsed.timeout_in_ms,
             context: parsed.context,
+            http_request: request,
             start_time_in_ms: start_time_in_ms
           )
 
@@ -155,17 +156,18 @@ def with_timeout(request, max_timeout_in_ms:)
         yield [max_timeout_in_ms, requested_timeout_in_ms].compact.min
       end
 
-      # Responsible for determining any `context` values to pass down into the `query_executor`,
-      # which in turn will make the values available to the GraphQL resolvers.
+      # Responsible for determining any extra `context` values to pass down into the `query_executor`,
+      # which in turn will make the values available to the GraphQL resolvers. (`http_request` itself is
+      # always passed down separately--see `process`--so it does not need to be added here.)
       #
-      # By default, our only context value is the HTTP request. This method exists to provide an extension
+      # By default we have no extra context values. This method exists to provide an extension
       # point so that ElasticGraph extensions can add `context` values based on the `request` as desired.
       #
       # Extensions can return an `HTTPResponse` with an error if the `request` is invalid according
       # to their requirements. Otherwise, they must call `super` (to delegate to this and any other
       # extensions) with a block. In the block, they must merge in their `context` values and then `yield`.
       def with_context(request)
-        yield({http_request: request})
+        yield({})
       end
 
       ParsedRequest = Data.define(:query_string, :variables, :operation_name, :timeout_in_ms, :context)

elasticgraph-graphql/lib/elastic_graph/graphql/query_adapter/requested_fields.rb

@@ -21,7 +21,7 @@ class WithoutSchema
           def call(field:, query:, lookahead:, args:, context:)
             return query if field.type.unwrap_fully.indexed_aggregation?
 
-            RequestedFields.new(context.fetch(:elastic_graph_schema)).call(
+            RequestedFields.new(context.elastic_graph_schema).call(
               field: field,
               query: query,
               lookahead: lookahead,

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

@@ -0,0 +1,78 @@
+# 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/graphql/client"
+require "elastic_graph/graphql/query_details_tracker"
+require "graphql"
+
+module ElasticGraph
+  class GraphQL
+    # `GraphQL::Query::Context` subclass used by ElasticGraph, providing typed accessors in
+    # place of untyped context hash keys. Registered as the `context_class` on the `GraphQL::Schema`
+    # built by `Schema`, via `new_class`, so that every query gets an instance closed over the
+    # `Schema` and `DatastoreSearchRouter` in use for that schema.
+    class QueryContext < ::GraphQL::Query::Context
+      def self.elastic_graph_schema
+        @elastic_graph_schema
+      end
+
+      def self.datastore_search_router
+        @datastore_search_router
+      end
+
+      # Builds a `QueryContext` subclass closed over the given `elastic_graph_schema` and
+      # `datastore_search_router`, suitable for registering as a `GraphQL::Schema#context_class`.
+      def self.new_class(elastic_graph_schema:, datastore_search_router:)
+        klass = ::Class.new(self) # : singleton(QueryContext)
+        klass.instance_variable_set(:@elastic_graph_schema, elastic_graph_schema)
+        klass.instance_variable_set(:@datastore_search_router, datastore_search_router)
+        klass
+      end
+
+      # `elastic_graph_schema` is only `nil` before `new_class` has closed over it, which never
+      # happens for a `QueryContext` actually used to execute a query (see `Schema#initialize`).
+      def elastic_graph_schema
+        self.class.elastic_graph_schema # : Schema
+      end
+
+      def datastore_search_router
+        self.class.datastore_search_router # : DatastoreSearchRouter
+      end
+
+      # @dynamic monotonic_clock_deadline, elastic_graph_client, http_request
+      attr_reader :monotonic_clock_deadline, :elastic_graph_client, :http_request
+
+      # Registers the given ElasticGraph-managed values on this context. Intended to be called
+      # exactly once, by `Schema#new_graphql_query` callers immediately after building a query and
+      # before making it available to resolvers or other user-facing code. Not for use elsewhere.
+      def register_elastic_graph_values(monotonic_clock_deadline: nil, elastic_graph_client: Client::ANONYMOUS, http_request: nil)
+        @monotonic_clock_deadline = monotonic_clock_deadline
+        @elastic_graph_client = elastic_graph_client
+        @http_request = http_request
+      end
+
+      # Lazily builds (and memoizes) the `QueryDetailsTracker` for this query.
+      def elastic_graph_query_tracker
+        @elastic_graph_query_tracker ||= QueryDetailsTracker.empty
+      end
+
+      # Caches the result of the given block (a built `DatastoreQuery`) under `cache_key`, scoped to
+      # this single query execution. See `Resolvers::QueryAdapter#build_query_from` for why this
+      # memoization is beneficial.
+      def cache_datastore_query(cache_key)
+        datastore_query_cache[cache_key] ||= yield
+      end
+
+      private
+
+      def datastore_query_cache
+        @datastore_query_cache ||= {}
+      end
+    end
+  end
+end

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

@@ -7,7 +7,6 @@
 # frozen_string_literal: true
 
 require "elastic_graph/graphql/client"
-require "elastic_graph/graphql/query_details_tracker"
 require "elastic_graph/support/hash_util"
 
 module ElasticGraph
@@ -40,6 +39,7 @@ def execute(
         timeout_in_ms: nil,
         operation_name: nil,
         context: {},
+        http_request: nil,
         start_time_in_ms: @monotonic_clock.now_in_ms
       )
         # Before executing the query, prune any null-valued variable fields. This means we
@@ -50,20 +50,18 @@ def execute(
         # due to a null-valued field referencing an undefined schema element.
         variables = ElasticGraph::Support::HashUtil.recursively_prune_nils_from(variables)
 
-        query_tracker = QueryDetailsTracker.empty
-
         query, result = build_and_execute_query(
           query_string: query_string,
           variables: variables,
           operation_name: operation_name,
           client: client,
-          context: context.merge({
-            monotonic_clock_deadline: timeout_in_ms&.+(start_time_in_ms),
-            elastic_graph_query_tracker: query_tracker,
-            elastic_graph_client: client
-          }.compact)
+          context: context,
+          monotonic_clock_deadline: timeout_in_ms&.+(start_time_in_ms),
+          http_request: http_request
         )
 
+        query_tracker = query.context.elastic_graph_query_tracker
+
         unless result.to_h.fetch("errors", []).empty?
           @logger.error <<~EOS
             Query #{query.selected_operation_name}[1] for client #{client.description} resulted in errors[2].
@@ -128,13 +126,18 @@ def execute(
 
       # Note: this is designed so that `elasticgraph-query_registry` can hook into this method. It needs to be able
       # to override how the query is built and executed.
-      def build_and_execute_query(query_string:, variables:, operation_name:, context:, client:)
+      def build_and_execute_query(query_string:, variables:, operation_name:, context:, client:, monotonic_clock_deadline:, http_request:)
         query = @schema.new_graphql_query(
           query_string,
           variables: variables,
           operation_name: operation_name,
           context: context
         )
+        query.context.register_elastic_graph_values(
+          monotonic_clock_deadline: monotonic_clock_deadline,
+          elastic_graph_client: client,
+          http_request: http_request
+        )
 
         [query, execute_query(query, client: client)]
       end

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

@@ -54,7 +54,7 @@ def object_type_hash
                 resolver_lambda =
                   if resolver.method(:resolve).parameters.include?([:keyreq, :lookahead])
                     lambda do |object, args, context|
-                      schema_field = context.fetch(:elastic_graph_schema).field_named(type_name, field_name)
+                      schema_field = context.elastic_graph_schema.field_named(type_name, field_name)
 
                       # Extract the `:lookahead` extra that we have configured all fields to provide.
                       # See https://graphql-ruby.org/api-doc/1.10.8/GraphQL/Execution/Lookahead.html for more info.
@@ -80,7 +80,7 @@ def object_type_hash
                     end
                   else
                     lambda do |object, args, context|
-                      schema_field = context.fetch(:elastic_graph_schema).field_named(type_name, field_name)
+                      schema_field = context.elastic_graph_schema.field_named(type_name, field_name)
                       # Convert args to the form they were defined in the schema, undoing the normalization
                       # the GraphQL gem does to convert them to Ruby keyword args form.
                       args = schema_field.args_to_schema_form(args)
@@ -109,7 +109,7 @@ def object_type_hash
 
         # In order to support unions and interfaces, we must implement `resolve_type`.
         def resolve_type(supertype, object, context)
-          schema = context.fetch(:elastic_graph_schema)
+          schema = context.elastic_graph_schema
           # If `__typename` is available, use that to resolve. It will be present on embedded abstract
           # types, and also on root documents indexed in a shared interface/union index.
           # (See `Inventor` in `config/schema/widgets.rb` for an example of an embedded abstract type.)
@@ -124,7 +124,8 @@ def resolve_type(supertype, object, context)
             # (See `Part` in `config/schema/widgets.rb` for an example of this kind of type union.)
             # This branch is only reached for individually-indexed types (no `__typename`
             # in the document), so the set always contains exactly one type.
-            schema.document_types_stored_in(object.index_definition_name).first.graphql_type
+            type = schema.document_types_stored_in(object.index_definition_name).first # : Schema::Type
+            type.graphql_type
           end
         end
       end