Add extension_data to QueryDetailsTracker for query registry logging

Add mechanism for GraphQL extensions to log additional data in
ElasticGraphQueryExecutorQueryDuration message. Query registry uses
this to log query_registration_status with values:
- matched_registered_query: exact match of registered query
- differing_registered_query: same operation name but query differs
- unregistered_query: client registered, no query with that name
- unregistered_client: client not registered at all

Key changes:
- Add extension_data field and []= method to QueryDetailsTracker
- Merge extension_data into QueryExecutor log output
- Update query validators to return registration_status as 3rd element
- Update GraphQL extension to set query_registration_status in tracker
- Optimize ForRegisteredClient to check operation name before
  canonicalization for better performance
- Use selected_operation_name instead of operation_name to correctly
  identify single-operation queries

Co-Authored-By: Claude <noreply@anthropic.com>

Author: myronmarston

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

@@ -17,6 +17,7 @@ class QueryDetailsTracker < Struct.new(
       :datastore_query_server_duration_ms,
       :datastore_query_client_duration_ms,
       :queried_shard_count,
+      :extension_data,
       :mutex
     )
       def self.empty
@@ -27,6 +28,7 @@ def self.empty
           datastore_query_server_duration_ms: 0,
           datastore_query_client_duration_ms: 0,
           queried_shard_count: 0,
+          extension_data: {},
           mutex: ::Thread::Mutex.new
         )
       end
@@ -52,6 +54,13 @@ def record_datastore_query_metrics(client_duration_ms:, server_duration_ms:, que
       def datastore_request_transport_duration_ms
         datastore_query_client_duration_ms - datastore_query_server_duration_ms
       end
+
+      # Allows extensions to set custom data that will be included in the query duration log.
+      def []=(key, value)
+        mutex.synchronize do
+          extension_data[key] = value
+        end
+      end
     end
   end
 end

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

@@ -117,7 +117,7 @@ def execute(
             "datastore_query_count" => query_tracker.query_counts_per_datastore_request.sum,
             "over_slow_threshold" => (duration > @slow_query_threshold_ms).to_s,
             "slo_result" => slo_result_for(query, duration)
-          })
+          }.merge(query_tracker.extension_data))
         end
 
         result

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

@@ -7,6 +7,7 @@ module ElasticGraph
       attr_accessor datastore_query_server_duration_ms: ::Integer
       attr_accessor datastore_query_client_duration_ms: ::Integer
       attr_accessor queried_shard_count: ::Integer
+      attr_accessor extension_data: ::Hash[::String, untyped]
       attr_accessor mutex: ::Thread::Mutex
 
       def initialize: (
@@ -16,6 +17,7 @@ module ElasticGraph
         datastore_query_server_duration_ms: ::Integer,
         datastore_query_client_duration_ms: ::Integer,
         queried_shard_count: ::Integer,
+        extension_data: ::Hash[::String, untyped],
         mutex: ::Thread::Mutex
       ) -> void
     end
@@ -30,6 +32,7 @@ module ElasticGraph
       ) -> void
 
       def datastore_request_transport_duration_ms: () -> ::Integer
+      def []=: (::String, untyped) -> untyped
     end
   end
 end

elasticgraph-graphql/sig/graphql_gem.rbs

@@ -113,6 +113,7 @@ module GraphQL
     attr_reader query_string: ::String?
 
     def selected_operation: () -> Language::Nodes::OperationDefinition?
+    def selected_operation_name: () -> ::String?
     def static_errors: () -> ::Array[_ValidationError]
 
     class Context

elasticgraph-graphql/spec/unit/elastic_graph/graphql/query_details_tracker_spec.rb

@@ -0,0 +1,51 @@
+# 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/query_details_tracker"
+
+module ElasticGraph
+  class GraphQL
+    RSpec.describe QueryDetailsTracker do
+      describe "#[]=" do
+        let(:tracker) { QueryDetailsTracker.empty }
+
+        it "allows extensions to set custom data in extension_data" do
+          tracker["custom_key"] = "custom_value"
+          expect(tracker.extension_data).to eq("custom_key" => "custom_value")
+        end
+
+        it "allows multiple values to be set" do
+          tracker["key1"] = "value1"
+          tracker["key2"] = "value2"
+
+          expect(tracker.extension_data).to eq(
+            "key1" => "value1",
+            "key2" => "value2"
+          )
+        end
+
+        it "allows overwriting existing extension data" do
+          tracker["key"] = "original_value"
+          tracker["key"] = "new_value"
+
+          expect(tracker.extension_data).to eq("key" => "new_value")
+        end
+
+        it "allows non-string values (per RBS signature)" do
+          tracker["int_key"] = 42
+          tracker["array_key"] = [1, 2, 3]
+
+          expect(tracker.extension_data).to eq(
+            "int_key" => 42,
+            "array_key" => [1, 2, 3]
+          )
+        end
+      end
+    end
+  end
+end

elasticgraph-graphql/spec/unit/elastic_graph/graphql/query_executor_spec.rb

@@ -369,6 +369,31 @@ class GraphQL
           }.to log(a_string_including("resulted in errors"))
         end
 
+        context "when extensions provide additional log data" do
+          it "includes extension data in the logged duration message" do
+            # Simulate an extension setting data in the query tracker
+            allow(::GraphQL::Execution::Interpreter).to receive(:run_all).and_wrap_original do |original, schema, queries, context:|
+              query_tracker = context[:elastic_graph_query_tracker]
+              query_tracker["custom_field"] = "custom_value"
+              query_tracker["another_field"] = "another_value"
+              original.call(schema, queries, context: context)
+            end
+
+            execute_expecting_no_errors(<<-QUERY, client: Client.new(name: "client-name", source_description: "client-description"))
+              query GetColors {
+                colors(args: {red: 12}) {
+                  red
+                }
+              }
+            QUERY
+
+            expect(logged_duration_message).to include(
+              "custom_field" => "custom_value",
+              "another_field" => "another_value"
+            )
+          end
+        end
+
         context "when the schema has been customized (as in an extension like elasticgraph-apollo)" do
           before(:context) do
             multiply_resolver = Class.new do

elasticgraph-query_registry/lib/elastic_graph/query_registry/client_data.rb

@@ -54,12 +54,18 @@ def with_updated_last_query(query_string, query)
       end
 
       def unregistered_query_error_for(query, client)
-        if operation_names.include?(query.operation_name.to_s)
-          "Query #{fingerprint_for(query)} differs from the registered form of `#{query.operation_name}` " \
+        # Note: we use `selected_operation_name` instead of `operation_name` because `operation_name` can return
+        # `nil` for single-operation queries when no explicit operation_name parameter is passed if accessed before
+        # the query AST is parsed, whereas `selected_operation_name`  parses the query AST and returns the operation
+        # name from the query document in that case.
+        selected_op_name = query.selected_operation_name.to_s
+
+        if operation_names.include?(selected_op_name)
+          "Query #{fingerprint_for(query)} differs from the registered form of `#{selected_op_name}` " \
           "for client #{client.description}."
         else
           "Query #{fingerprint_for(query)} is unregistered; client #{client.description} has no " \
-          "registered query with a `#{query.operation_name}` operation."
+          "registered query with a `#{selected_op_name}` operation."
         end
       end
 

elasticgraph-query_registry/lib/elastic_graph/query_registry/graphql_extension.rb

@@ -60,14 +60,16 @@ def initialize(
       private
 
       def build_and_execute_query(query_string:, variables:, operation_name:, context:, client:)
-        query, errors = @registry.build_and_validate_query(
+        query, errors, registration_status = @registry.build_and_validate_query(
           query_string,
           variables: variables,
           operation_name: operation_name,
           context: context,
           client: client
         )
 
+        context.fetch(:elastic_graph_query_tracker)["query_registration_status"] = registration_status
+
         if errors.empty?
           [query, execute_query(query, client: client)]
         else

elasticgraph-query_registry/lib/elastic_graph/query_registry/query_validators/for_registered_client.rb

@@ -8,6 +8,7 @@
 
 require "elastic_graph/graphql/client"
 require "elastic_graph/query_registry/client_data"
+require "elastic_graph/query_registry/registration_status"
 require "graphql"
 
 module ElasticGraph
@@ -44,19 +45,34 @@ def build_and_validate_query(query_string, client:, variables: {}, operation_nam
 
           if (cached_query = client_data.cached_query_for(query_string.to_s))
             prepared_query = prepare_query_for_execution(cached_query, variables: variables, operation_name: operation_name, context: context)
-            return [prepared_query, []]
+            return [prepared_query, [], RegistrationStatus::MATCHED_REGISTERED_QUERY]
           end
 
           query = yield
 
+          # Check operation name first (fast O(1) set lookup) to avoid canonicalization when not needed.
+          # Note: we use `selected_operation_name` instead of `operation_name` because `operation_name` returns
+          # `nil` for single-operation queries when no explicit operation_name parameter is passed, whereas
+          # `selected_operation_name` returns the operation name from the query document in that case.
+          registration_status =
+            if client_data.operation_names.include?(query.selected_operation_name.to_s)
+              if client_data.canonical_query_strings.include?(ClientData.canonical_query_string_from(query, schema_element_names: schema.element_names))
+                RegistrationStatus::MATCHED_REGISTERED_QUERY
+              else
+                RegistrationStatus::DIFFERING_REGISTERED_QUERY
+              end
+            else
+              RegistrationStatus::UNREGISTERED_QUERY
+            end
+
           # This client allows any query, so we can just return the query with no errors here.
           # Note: we could put this at the top of the method, but if the query is registered and matches
           # the registered form, the `cached_query` above is more efficient as it avoids unnecessarily
           # parsing the query.
-          return [query, []] if allow_any_query_for_clients.include?(client.name)
+          return [query, [], registration_status] if allow_any_query_for_clients.include?(client.name)
 
-          if !client_data.canonical_query_strings.include?(ClientData.canonical_query_string_from(query, schema_element_names: schema.element_names))
-            return [query, [client_data.unregistered_query_error_for(query, client)]]
+          unless registration_status == RegistrationStatus::MATCHED_REGISTERED_QUERY
+            return [query, [client_data.unregistered_query_error_for(query, client)], registration_status]
           end
 
           # The query is slightly different from a registered query, but not in any material fashion
@@ -74,7 +90,7 @@ def build_and_validate_query(query_string, client:, variables: {}, operation_nam
             (_ = cached_client_data).with_updated_last_query(query_string, cachable_query)
           end
 
-          [query, []]
+          [query, [], registration_status]
         end
 
         private

elasticgraph-query_registry/lib/elastic_graph/query_registry/query_validators/for_unregistered_client.rb

@@ -6,6 +6,8 @@
 #
 # frozen_string_literal: true
 
+require "elastic_graph/query_registry/registration_status"
+
 module ElasticGraph
   module QueryRegistry
     module QueryValidators
@@ -15,15 +17,15 @@ module QueryValidators
         def build_and_validate_query(query_string, client:, variables: {}, operation_name: nil, context: {})
           query = yield
 
-          return [query, []] if allow_unregistered_clients
+          return [query, [], RegistrationStatus::UNREGISTERED_CLIENT] if allow_unregistered_clients
 
           client_name = client&.name
-          return [query, []] if client_name && allow_any_query_for_clients.include?(client_name)
+          return [query, [], RegistrationStatus::UNREGISTERED_CLIENT] if client_name && allow_any_query_for_clients.include?(client_name)
 
           [query, [
             "Client #{client&.description || "(unknown)"} is not a registered client, it is not in " \
             "`allow_any_query_for_clients` and `allow_unregistered_clients` is false."
-          ]]
+          ], RegistrationStatus::UNREGISTERED_CLIENT]
         end
       end
     end