Add query execution context based tracing for API requests

lib/graphql/tracing.rb

@@ -8,6 +8,7 @@
 require "graphql/tracing/scout_tracing"
 require "graphql/tracing/statsd_tracing"
 require "graphql/tracing/prometheus_tracing"
+require "graphql/tracing/opentelemetry_tracing"
 
 if defined?(PrometheusExporter::Server)
   require "graphql/tracing/prometheus_tracing/graphql_collector"

lib/graphql/tracing/opentelemetry_tracing.rb

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+module GraphQL
+  module Tracing
+    class OpenTelemetryTracing < PlatformTracing
+      self.platform_keys = {
+        'lex' => 'graphql.lex',
+        'parse' => 'graphql.parse',
+        'validate' => 'graphql.validate',
+        'analyze_query' => 'graphql.analyze_query',
+        'analyze_multiplex' => 'graphql.analyze_multiplex',
+        'execute_query' => 'graphql.execute_query',
+        'execute_query_lazy' => 'graphql.execute_query_lazy',
+        'execute_multiplex' => 'graphql.execute_multiplex'
+      }
+
+      def platform_trace(platform_key, key, data)
+        return yield if platform_key.nil?
+
+        tracer.in_span(platform_key, attributes: attributes_for(key, data)) do |span, _context|
+          yield.tap do |response|
+            errors = response[:errors]&.compact&.map { |e| e.to_h }&.to_json if key == 'validate'
+            unless errors.nil?
+              span.add_event(
+                'graphql.validation.error',
+                attributes: {
+                  'message' => errors
+                }
+              )
+            end
+          end
+        end
+      end
+
+      def platform_field_key(type, field)
+        "#{type.graphql_name}.#{field.graphql_name}"
+      end
+
+      def platform_authorized_key(type)
+        "#{type.graphql_name}.authorized"
+      end
+
+      def platform_resolve_type_key(type)
+        "#{type.graphql_name}.resolve_type"
+      end
+
+      private
+
+      def tracer
+        OpenTelemetry::Instrumentation::GraphQL::Instrumentation.instance.tracer
+      end
+
+      def config
+        OpenTelemetry::Instrumentation::GraphQL::Instrumentation.instance.config
+      end
+
+      def platform_key_enabled?(ctx, key)
+        return false unless config[key]
+
+        ns = ctx.namespace(:opentelemetry)
+        return true if ns.empty? # restores original behavior so that keys are returned if tracing is not set in context.
+        return false unless ns.key?(key) && ns[key]
+
+        return true
+      end
+
+      def attributes_for(key, data)
+        attributes = {}
+        case key
+        when 'execute_query'
+          attributes['selected_operation_name'] = data[:query].selected_operation_name if data[:query].selected_operation_name
+          attributes['selected_operation_type'] = data[:query].selected_operation.operation_type
+          attributes['query_string'] = data[:query].query_string
+        end
+        attributes
+      end
+
+      def cached_platform_key(ctx, key, trace_phase)
+        cache = ctx.namespace(self.class)[:platform_key_cache] ||= {}
+
+        cache.fetch(key) do
+          cache[key] = if trace_phase == :field
+            return unless platform_key_enabled?(ctx, :enable_platform_field)
+
+            yield
+          elsif trace_phase == :authorized
+            return unless platform_key_enabled?(ctx, :enable_platform_authorized)
+
+            yield
+          elsif trace_phase == :resolve_type
+            return unless platform_key_enabled?(ctx, :enable_platform_resolve_type)
+
+            yield
+          else
+            raise "Unknown trace phase"
+          end
+        end
+      end
+    end
+  end
+end

lib/graphql/tracing/platform_tracing.rb

@@ -45,7 +45,7 @@ def trace(key, data)
 
             platform_key = if trace_field
               context = data.fetch(:query).context
-              cached_platform_key(context, field) { platform_field_key(data[:owner], field) }
+              cached_platform_key(context, field, :field) { platform_field_key(data[:owner], field) }
             else
               nil
             end
@@ -61,14 +61,14 @@ def trace(key, data)
         when "authorized", "authorized_lazy"
           type = data.fetch(:type)
           context = data.fetch(:context)
-          platform_key = cached_platform_key(context, type) { platform_authorized_key(type) }
+          platform_key = cached_platform_key(context, type, :authorized) { platform_authorized_key(type) }
           platform_trace(platform_key, key, data) do
             yield
           end
         when "resolve_type", "resolve_type_lazy"
           type = data.fetch(:type)
           context = data.fetch(:context)
-          platform_key = cached_platform_key(context, type) { platform_resolve_type_key(type) }
+          platform_key = cached_platform_key(context, type, :resolve_type) { platform_resolve_type_key(type) }
           platform_trace(platform_key, key, data) do
             yield
           end
@@ -116,7 +116,7 @@ def fallback_transaction_name(context)
       # If the key isn't present, the given block is called and the result is cached for `key`.
       #
       # @return [String]
-      def cached_platform_key(ctx, key)
+      def cached_platform_key(ctx, key, trace_phase)
         cache = ctx.namespace(self.class)[:platform_key_cache] ||= {}
         cache.fetch(key) { cache[key] = yield }
       end

spec/graphql/tracing/opentelemetry_tracing_spec.rb

@@ -0,0 +1,156 @@
+# frozen_string_literal: true
+require "spec_helper"
+
+describe ::GraphQL::Tracing::OpenTelemetryTracing do
+  module OpenTelemetryTest
+    class Thing < GraphQL::Schema::Object
+      implements GraphQL::Types::Relay::Node
+    end
+
+    class Query < GraphQL::Schema::Object
+      include GraphQL::Types::Relay::HasNodeField
+
+      field :int, Integer, null: false
+
+      def int
+        1
+      end
+    end
+
+    class SchemaWithPerRequestTracing < GraphQL::Schema
+      query(Query)
+      use(GraphQL::Tracing::OpenTelemetryTracing)
+      orphan_types(Thing)
+
+      def self.object_from_id(_id, _ctx)
+        :thing
+      end
+
+      def self.resolve_type(_type, _obj, _ctx)
+        Thing
+      end
+    end
+  end
+
+  before do
+    OpenTelemetry::Instrumentation::GraphQL::Instrumentation.clear_all
+  end
+
+  it "captures all keys when tracing is enabled in config and in query execution context" do
+    query = GraphQL::Query.new(OpenTelemetryTest::SchemaWithPerRequestTracing, '{ node(id: "1") { __typename } }')
+    query.context.namespace(:opentelemetry)[:enable_platform_field] = true
+    query.context.namespace(:opentelemetry)[:enable_platform_authorized] = true
+    query.context.namespace(:opentelemetry)[:enable_platform_resolve_type] = true
+
+    query.result
+
+    assert_span("Query.authorized")
+    assert_span("Query.node")
+    assert_span("Node.resolve_type")
+    assert_span("Thing.authorized")
+    assert_span("DynamicFields.authorized")
+  end
+
+  it "does not capture the keys when tracing is not enabled in config but is enabled in query execution context" do
+    OpenTelemetry::Instrumentation::GraphQL::Instrumentation.instance.stub(:config, {
+      schemas: [],
+      enable_platform_field: false,
+      enable_platform_authorized: false,
+      enable_platform_resolve_type: false
+    }) do
+
+      query = GraphQL::Query.new(OpenTelemetryTest::SchemaWithPerRequestTracing, '{ node(id: "1") { __typename } }')
+      query.context.namespace(:opentelemetry)[:enable_platform_field] = true
+      query.context.namespace(:opentelemetry)[:enable_platform_authorized] = true
+      query.context.namespace(:opentelemetry)[:enable_platform_resolve_type] = true
+
+      query.result
+
+      refute_span("Query.authorized")
+      refute_span("Query.node")
+      refute_span("Node.resolve_type")
+      refute_span("Thing.authorized")
+      refute_span("DynamicFields.authorized")
+    end
+  end
+
+  it "does not capture any key when tracing is not enabled in config and tracing is not set in context" do
+    OpenTelemetry::Instrumentation::GraphQL::Instrumentation.instance.stub(:config, {
+      schemas: [],
+      enable_platform_field: false,
+      enable_platform_authorized: false,
+      enable_platform_resolve_type: false
+    }) do
+
+      query = GraphQL::Query.new(OpenTelemetryTest::SchemaWithPerRequestTracing, '{ node(id: "1") { __typename } }')
+
+      query.result
+
+      refute_span("Query.authorized")
+      refute_span("Query.node")
+      refute_span("Node.resolve_type")
+      refute_span("Thing.authorized")
+      refute_span("DynamicFields.authorized")
+    end
+  end
+
+  it "does not capture any key when tracing is not enabled in config and context" do
+    OpenTelemetry::Instrumentation::GraphQL::Instrumentation.instance.stub(:config, {
+      schemas: [],
+      enable_platform_field: false,
+      enable_platform_authorized: false,
+      enable_platform_resolve_type: false
+    }) do
+
+      query = GraphQL::Query.new(OpenTelemetryTest::SchemaWithPerRequestTracing, '{ node(id: "1") { __typename } }')
+      query.context.namespace(:opentelemetry)[:enable_platform_field] = false
+      query.context.namespace(:opentelemetry)[:enable_platform_authorized] = false
+      query.context.namespace(:opentelemetry)[:enable_platform_resolve_type] = false
+
+      query.result
+
+      refute_span("Query.authorized")
+      refute_span("Query.node")
+      refute_span("Node.resolve_type")
+      refute_span("Thing.authorized")
+      refute_span("DynamicFields.authorized")
+    end
+  end
+
+  it "captures all keys when tracing is enabled in config but is not set in context" do
+    query = GraphQL::Query.new(OpenTelemetryTest::SchemaWithPerRequestTracing, '{ node(id: "1") { __typename } }')
+
+    query.result
+
+    assert_span("Query.authorized")
+    assert_span("Query.node")
+    assert_span("Node.resolve_type")
+    assert_span("Thing.authorized")
+    assert_span("DynamicFields.authorized")
+  end
+
+  it "does not capture any key when tracing is enabled in config but is not enabled in context" do
+    query = GraphQL::Query.new(OpenTelemetryTest::SchemaWithPerRequestTracing, '{ node(id: "1") { __typename } }')
+    query.context.namespace(:opentelemetry)[:enable_platform_field] = false
+    query.context.namespace(:opentelemetry)[:enable_platform_authorized] = false
+    query.context.namespace(:opentelemetry)[:enable_platform_resolve_type] = false
+
+    query.result
+
+    refute_span("Query.authorized")
+    refute_span("Query.node")
+    refute_span("Node.resolve_type")
+    refute_span("Thing.authorized")
+    refute_span("DynamicFields.authorized")
+  end
+
+  private
+
+  def assert_span(span)
+    assert OpenTelemetry::Instrumentation::GraphQL::Instrumentation::EVENTS.include?(span)
+  end
+
+  def refute_span(span)
+    refute OpenTelemetry::Instrumentation::GraphQL::Instrumentation::EVENTS.include?(span)
+  end
+end

spec/support/opentelemetry.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+# A stub for the Opentelemetry agent, so we can make assertions about how it is used
+if defined?(OpenTelemetry)
+  raise "Expected Opentelemetry to be undefined, so that we could define a stub for it."
+end
+
+module OpenTelemetry
+  module Instrumentation
+    module GraphQL
+      class Instrumentation
+        EVENTS = []
+        class << self
+          def instance
+            @instance ||= new
+          end
+
+          def clear_all
+            EVENTS.clear
+          end
+        end
+
+        def tracer
+          @tracer ||= DummyTracer.new
+        end
+
+        def config
+          @config ||= {
+            schemas: [],
+            enable_platform_field: true,
+            enable_platform_authorized: true,
+            enable_platform_resolve_type: true
+          }
+        end
+      end
+
+      class DummyTracer
+        class TestSpan
+          def add_event(name, attributes:)
+            self
+          end
+        end
+
+        def in_span(name, attributes: nil, links: nil, start_timestamp: nil, kind: nil)
+          OpenTelemetry::Instrumentation::GraphQL::Instrumentation::EVENTS << name
+
+          yield(TestSpan.new, {})
+        end
+      end
+    end
+  end
+end