Skip to content

Add extension_data to QueryDetailsTracker for query registry logging #981

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
myronmarston merged 1 commit into from
Jan 23, 2026
Merged

Add extension_data to QueryDetailsTracker for query registry logging #981

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions elasticgraph-graphql/lib/elastic_graph/graphql/query_details_tracker.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
2 changes: 1 addition & 1 deletion elasticgraph-graphql/lib/elastic_graph/graphql/query_executor.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions elasticgraph-graphql/sig/elastic_graph/graphql/query_details_tracker.rbs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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: (
Expand All @@ -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
Expand All @@ -30,6 +32,7 @@ module ElasticGraph
) -> void

def datastore_request_transport_duration_ms: () -> ::Integer
def []=: (::String, untyped) -> untyped
end
end
end
1 change: 1 addition & 0 deletions elasticgraph-graphql/sig/graphql_gem.rbs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
51 changes: 51 additions & 0 deletions elasticgraph-graphql/spec/unit/elastic_graph/graphql/query_details_tracker_spec.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
25 changes: 25 additions & 0 deletions elasticgraph-graphql/spec/unit/elastic_graph/graphql/query_executor_spec.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 9 additions & 3 deletions elasticgraph-query_registry/lib/elastic_graph/query_registry/client_data.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
4 changes: 3 additions & 1 deletion elasticgraph-query_registry/lib/elastic_graph/query_registry/graphql_extension.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
26 changes: 21 additions & 5 deletions ...query_registry/lib/elastic_graph/query_registry/query_validators/for_registered_client.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
8 changes: 5 additions & 3 deletions ...ery_registry/lib/elastic_graph/query_registry/query_validators/for_unregistered_client.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
#
# frozen_string_literal: true

require "elastic_graph/query_registry/registration_status"

module ElasticGraph
module QueryRegistry
module QueryValidators
Expand All @@ -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
Expand Down
26 changes: 26 additions & 0 deletions elasticgraph-query_registry/lib/elastic_graph/query_registry/registration_status.rb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
# 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

module ElasticGraph
module QueryRegistry
# Constants for query registration status values logged in `ElasticGraphQueryExecutorQueryDuration`.
module RegistrationStatus
# Query exactly matched a registered query (used from cache).
MATCHED_REGISTERED_QUERY = "matched_registered_query"

# Query has same operation name as a registered query but query body differs.
DIFFERING_REGISTERED_QUERY = "differing_registered_query"

# Client is registered but has no query with this operation name.
UNREGISTERED_QUERY = "unregistered_query"

# Client is not registered in the query registry.
UNREGISTERED_CLIENT = "unregistered_client"
end
end
end
2 changes: 1 addition & 1 deletion ...uery_registry/sig/elastic_graph/query_registry/query_validators/for_registered_client.rbs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ module ElasticGraph
?variables: ::Hash[::String, untyped],
?operation_name: ::String?,
?context: ::Hash[::Symbol, untyped]
) { () -> ::GraphQL::Query } -> [::GraphQL::Query, ::Array[::String]]
) { () -> ::GraphQL::Query } -> [::GraphQL::Query, ::Array[::String], ::String]

private

Expand Down
2 changes: 1 addition & 1 deletion ...ry_registry/sig/elastic_graph/query_registry/query_validators/for_unregistered_client.rbs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ module ElasticGraph
?variables: ::Hash[::String, untyped],
?operation_name: ::String?,
?context: ::Hash[::Symbol, untyped]
) { () -> ::GraphQL::Query } -> [::GraphQL::Query, ::Array[::String]]
) { () -> ::GraphQL::Query } -> [::GraphQL::Query, ::Array[::String], ::String]
end
end
end
Expand Down
10 changes: 10 additions & 0 deletions elasticgraph-query_registry/sig/elastic_graph/query_registry/registration_status.rbs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
module ElasticGraph
module QueryRegistry
module RegistrationStatus
MATCHED_REGISTERED_QUERY: ::String
DIFFERING_REGISTERED_QUERY: ::String
UNREGISTERED_QUERY: ::String
UNREGISTERED_CLIENT: ::String
end
end
end
2 changes: 1 addition & 1 deletion elasticgraph-query_registry/sig/elastic_graph/query_registry/registry.rbs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ module ElasticGraph
?variables: ::Hash[::String, untyped],
?operation_name: ::String?,
?context: ::Hash[::Symbol, untyped]
) -> [::GraphQL::Query, ::Array[::String]]
) -> [::GraphQL::Query, ::Array[::String], ::String]

private

Expand Down
Loading