Add new GraphQL tracer complying to span attributes specification

Steepfile

@@ -595,6 +595,7 @@ target :datadog do
   library 'opentelemetry-api'
   library 'passenger'
   library 'webmock'
+  library 'graphql'
 
   # TODO: gem 'libddwaf'
   library 'libddwaf'

docs/GettingStarted.md

@@ -814,11 +814,12 @@ YourSchema.execute(query, variables: {}, context: {}, operation_name: nil)
 
 The `instrument :graphql` method accepts the following parameters. Additional options can be substituted in for `options`:
 
-| Key                      | Type     | Description                                                                                                                                                  | Default          |
-| ------------------------ | -------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------------- |
-| `schemas`                | `Array`  | Array of `GraphQL::Schema` objects (that support class-based schema only) to trace. If you do not provide any, then tracing will applied to all the schemas. | `[]`             |
-| `with_deprecated_tracer` | `Bool`   | Enable to instrument with deprecated `GraphQL::Tracing::DataDogTracing`. Default is `false`, using `GraphQL::Tracing::DataDogTrace`                          | `false`          |
-| `service_name`           | `String` | Service name used for graphql instrumentation                                                                                                                | `'ruby-graphql'` |
+| Key                      | Type     | Description                                                                                                                                                                                                     | Default          |
+| ------------------------ | -------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------- |
+| `schemas`                | `Array`  | Array of `GraphQL::Schema` objects (that support class-based schema only) to trace. If you do not provide any, then tracing will applied to all the schemas.                                                    | `[]`             |
+| `with_unified_tracer`    | `Bool`   | Enable to instrument with `UnifiedTrace` tracer, enabling support for API Catalog. `with_deprecated_tracer` has priority over this. Default is `false`, using `GraphQL::Tracing::DataDogTrace` (Added in v2.2)  | `false`          |
+| `with_deprecated_tracer` | `Bool`   | Enable to instrument with deprecated `GraphQL::Tracing::DataDogTracing`. This has priority over `with_unified_tracer`. Default is `false`, using `GraphQL::Tracing::DataDogTrace`                               | `false`          |
+| `service_name`           | `String` | Service name used for graphql instrumentation                                                                                                                                                                   | `'ruby-graphql'` |
 
 **Manually configuring GraphQL schemas**
 
@@ -832,6 +833,14 @@ class YourSchema < GraphQL::Schema
 end
 ```
 
+With `UnifiedTracer` (Added in v2.2)
+
+```ruby
+class YourSchema < GraphQL::Schema
+  trace_with Datadog::Tracing::Contrib::GraphQL::UnifiedTrace
+end
+```
+
 or with `GraphQL::Tracing::DataDogTracing` (deprecated)
 
 ```ruby

lib/datadog/tracing/contrib/graphql/configuration/settings.rb

@@ -42,6 +42,11 @@ class Settings < Contrib::Configuration::Settings
               o.type :bool
               o.default false
             end
+
+            option :with_unified_tracer do |o|
+              o.type :bool
+              o.default false
+            end
           end
         end
       end

lib/datadog/tracing/contrib/graphql/patcher.rb

@@ -4,6 +4,7 @@
 require_relative '../patcher'
 require_relative 'tracing_patcher'
 require_relative 'trace_patcher'
+require_relative 'unified_trace_patcher'
 
 module Datadog
   module Tracing
@@ -23,10 +24,15 @@ def patch
             if configuration[:with_deprecated_tracer]
               TracingPatcher.patch!(schemas, trace_options)
             elsif Integration.trace_supported?
-              TracePatcher.patch!(schemas, trace_options)
+              if configuration[:with_unified_tracer]
+                UnifiedTracePatcher.patch!(schemas, trace_options)
+              else
+                TracePatcher.patch!(schemas, trace_options)
+              end
             else
               Datadog.logger.warn(
-                "GraphQL version (#{target_version}) does not support GraphQL::Tracing::DataDogTrace. "\
+                "GraphQL version (#{target_version}) does not support GraphQL::Tracing::DataDogTrace"\
+                'or Datadog::Tracing::Contrib::GraphQL::UnifiedTrace.'\
                 'Falling back to GraphQL::Tracing::DataDogTracing.'
               )
               TracingPatcher.patch!(schemas, trace_options)

lib/datadog/tracing/contrib/graphql/unified_trace.rb

@@ -0,0 +1,170 @@
+# frozen_string_literal: true
+
+require 'graphql/tracing'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module GraphQL
+        # These methods will be called by the GraphQL runtime to trace the execution of queries
+        module UnifiedTrace
+          # @param analytics_enabled [Boolean] Deprecated
+          # @param analytics_sample_rate [Float] Deprecated
+          # @param service [String|nil] The service name to be set on the spans
+          def initialize(*args, analytics_enabled: false, analytics_sample_rate: 1.0, service: nil, **kwargs)
+            @analytics_enabled = analytics_enabled
+            @analytics_sample_rate = analytics_sample_rate
+
+            @service_name = service
+            @has_prepare_span = respond_to?(:prepare_span)
+            super
+          end
+
+          def lex(*args, query_string:, **kwargs)
+            trace(proc { super }, 'lex', query_string, query_string: query_string)
+          end
+
+          def parse(*args, query_string:, **kwargs)
+            trace(proc { super }, 'parse', query_string, query_string: query_string) do |span|
+              span.set_tag('graphql.source', query_string)
+            end
+          end
+
+          def validate(*args, query:, validate:, **kwargs)
+            trace(proc { super }, 'validate', query.selected_operation_name, query: query, validate: validate) do |span|
+              span.set_tag('graphql.source', query.query_string)
+            end
+          end
+
+          def analyze_multiplex(*args, multiplex:, **kwargs)
+            trace(proc { super }, 'analyze_multiplex', multiplex_resource(multiplex), multiplex: multiplex)
+          end
+
+          def analyze_query(*args, query:, **kwargs)
+            trace(proc { super }, 'analyze', query.query_string, query: query)
+          end
+
+          def execute_multiplex(*args, multiplex:, **kwargs)
+            trace(proc { super }, 'execute_multiplex', multiplex_resource(multiplex), multiplex: multiplex) do |span|
+              span.set_tag('graphql.source', "Multiplex[#{multiplex.queries.map(&:query_string).join(', ')}]")
+            end
+          end
+
+          def execute_query(*args, query:, **kwargs)
+            trace(proc { super }, 'execute', query.selected_operation_name, query: query) do |span|
+              span.set_tag('graphql.source', query.query_string)
+              span.set_tag('graphql.operation.type', query.selected_operation.operation_type)
+              span.set_tag('graphql.operation.name', query.selected_operation_name) if query.selected_operation_name
+              query.provided_variables.each do |key, value|
+                span.set_tag("graphql.variables.#{key}", value)
+              end
+            end
+          end
+
+          def execute_query_lazy(*args, query:, multiplex:, **kwargs)
+            resource = if query
+                         query.selected_operation_name || fallback_transaction_name(query.context)
+                       else
+                         multiplex_resource(multiplex)
+                       end
+            trace(proc { super }, 'execute_lazy', resource, query: query, multiplex: multiplex)
+          end
+
+          def execute_field_span(callable, span_key, **kwargs)
+            platform_key = @platform_key_cache[UnifiedTrace].platform_field_key_cache[kwargs[:field]]
+
+            if platform_key
+              trace(callable, span_key, platform_key, **kwargs) do |span|
+                kwargs[:arguments].each do |key, value|
+                  span.set_tag("graphql.variables.#{key}", value)
+                end
+              end
+            else
+              callable.call
+            end
+          end
+
+          def execute_field(*args, **kwargs)
+            # kwargs[:arguments] is { id => 1 } for 'user(id: 1) { name }'. This is what we want to send to the WAF.
+            execute_field_span(proc { super }, 'resolve', **kwargs)
+          end
+
+          def execute_field_lazy(*args, **kwargs)
+            execute_field_span(proc { super }, 'resolve_lazy', **kwargs)
+          end
+
+          def authorized_span(callable, span_key, **kwargs)
+            platform_key = @platform_key_cache[UnifiedTrace].platform_authorized_key_cache[kwargs[:type]]
+            trace(callable, span_key, platform_key, **kwargs)
+          end
+
+          def authorized(*args, **kwargs)
+            authorized_span(proc { super }, 'authorized', **kwargs)
+          end
+
+          def authorized_lazy(*args, **kwargs)
+            authorized_span(proc { super }, 'authorized_lazy', **kwargs)
+          end
+
+          def resolve_type_span(callable, span_key, **kwargs)
+            platform_key = @platform_key_cache[UnifiedTrace].platform_resolve_type_key_cache[kwargs[:type]]
+            trace(callable, span_key, platform_key, **kwargs)
+          end
+
+          def resolve_type(*args, **kwargs)
+            resolve_type_span(proc { super }, 'resolve_type', **kwargs)
+          end
+
+          def resolve_type_lazy(*args, **kwargs)
+            resolve_type_span(proc { super }, 'resolve_type_lazy', **kwargs)
+          end
+
+          include ::GraphQL::Tracing::PlatformTrace
+
+          # Implement this method in a subclass to apply custom tags to datadog spans
+          # @param key [String] The event being traced
+          # @param data [Hash] The runtime data for this event (@see GraphQL::Tracing for keys for each event)
+          # @param span [Datadog::Tracing::SpanOperation] The datadog span for this event
+          # def prepare_span(key, data, span)
+          # end
+
+          def platform_field_key(field, *args, **kwargs)
+            field.path
+          end
+
+          def platform_authorized_key(type, *args, **kwargs)
+            "#{type.graphql_name}.authorized"
+          end
+
+          def platform_resolve_type_key(type, *args, **kwargs)
+            "#{type.graphql_name}.resolve_type"
+          end
+
+          private
+
+          def trace(callable, trace_key, resource, **kwargs)
+            Tracing.trace("graphql.#{trace_key}", resource: resource, service: @service_name, type: 'graphql') do |span|
+              yield(span) if block_given?
+
+              prepare_span(trace_key, kwargs, span) if @has_prepare_span
+
+              callable.call
+            end
+          end
+
+          def multiplex_resource(multiplex)
+            return nil unless multiplex
+
+            operations = multiplex.queries.map(&:selected_operation_name).compact.join(', ')
+            if operations.empty?
+              first_query = multiplex.queries.first
+              fallback_transaction_name(first_query && first_query.context)
+            else
+              operations
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module GraphQL
+        # Provides instrumentation for `graphql` through the GraphQL's tracing with methods defined in UnifiedTrace
+        module UnifiedTracePatcher
+          module_function
+
+          def patch!(schemas, options)
+            require_relative 'unified_trace'
+            if schemas.empty?
+              ::GraphQL::Schema.trace_with(UnifiedTrace, **options)
+            else
+              schemas.each do |schema|
+                schema.trace_with(UnifiedTrace, **options)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

sig/datadog/tracing/contrib/graphql/unified_trace.rbs

@@ -0,0 +1,76 @@
+module Datadog
+  module Tracing
+    module Contrib
+      module GraphQL
+        module UnifiedTrace
+          @analytics_enabled: bool
+
+          @analytics_sample_rate: Float
+
+          @service_name: String?
+
+          @has_prepare_span: bool
+          def initialize: (?analytics_enabled: bool, ?analytics_sample_rate: Float, ?service: String?, **Hash[Symbol, Object] kwargs) -> self
+
+          type lexerArray = Array[Integer | Symbol | String | nil | lexerArray]
+
+          def lex: (query_string: String) -> lexerArray
+
+          def parse: (query_string: String) -> GraphQL::Language::Nodes::Document
+
+          def validate: (query: GraphQL::Query, validate: bool) -> { remaining_timeout: Float?, error: Array[StandardError] }
+
+          def analyze_multiplex: (multiplex: GraphQL::Execution::Multiplex) -> Array[Object]
+
+          def analyze_query: (query: GraphQL::Query) -> Array[Object]
+
+          def execute_multiplex: (multiplex: GraphQL::Execution::Multiplex) -> Array[GraphQL::Query::Result]
+
+          def execute_query: (query: GraphQL::Query) -> GraphQL::Query::Result
+
+          def execute_query_lazy: (query: GraphQL::Query, multiplex: GraphQL::Execution::Multiplex) -> GraphQL::Query::Result
+
+          type executeFieldKwargs = {query: GraphQL::Query, field: GraphQL::Schema::Field, ast_node: GraphQL::Language::Nodes::Field, arguments: Hash[Symbol, String], object: GraphQL::Schema::Object?}
+
+          def execute_field_span: (Proc callable, String span_key, **executeFieldKwargs kwargs) -> Array[Object]
+
+          def execute_field: (**executeFieldKwargs kwargs) -> Array[Object]
+
+          def execute_field_lazy: (**executeFieldKwargs kwargs) -> Array[Object]
+
+          type authorizedKwargs = {query: GraphQL::Query, type: GraphQL::Schema::Object, object: GraphQL::Schema::Object?}
+
+          def authorized_span: (Proc callable, String span_key, **authorizedKwargs kwargs) -> GraphQL::Schema::Object?
+
+          def authorized: (**authorizedKwargs kwargs) -> GraphQL::Schema::Object?
+
+          def authorized_lazy: (**authorizedKwargs kwargs) -> GraphQL::Schema::Object?
+
+          type resolveTypeKwargs = {query: GraphQL::Query, type: GraphQL::Schema::Union, object: GraphQL::Schema::Object?}
+
+          def resolve_type_span: (Proc callable, String span_key, **resolveTypeKwargs kwargs) -> [GraphQL::Schema::Object, nil]
+
+          def resolve_type: (**resolveTypeKwargs kwargs) -> [GraphQL::Schema::Object, nil]
+
+          def resolve_type_lazy: (**resolveTypeKwargs kwargs) -> [GraphQL::Schema::Object, nil]
+
+          def platform_field_key: (GraphQL::Schema::Field field) -> String
+
+          def platform_authorized_key: (GraphQL::Schema::Object type) -> String
+
+          def platform_resolve_type_key: (GraphQL::Schema::Union type) -> String
+
+          private
+
+          type traceKwargsValues = GraphQL::Query | GraphQL::Schema::Union | GraphQL::Schema::Object | GraphQL::Schema::Field | GraphQL::Execution::Multiplex | GraphQL::Language::Nodes::Field | Hash[Symbol, String] | String | bool | nil
+
+          type traceResult = lexerArray | GraphQL::Language::Nodes::Document | { remaining_timeout: Float?, error: Array[StandardError] } | Array[Object] | GraphQL::Schema::Object? | [GraphQL::Schema::Object, nil]
+
+          def trace: (Proc callable, String trace_key, String resource, **Hash[Symbol, traceKwargsValues ] kwargs) ?{ (Datadog::Tracing::SpanOperation) -> void } -> traceResult
+
+          def multiplex_resource: (GraphQL::Execution::Multiplex multiplex) -> String?
+        end
+      end
+    end
+  end
+end

sig/datadog/tracing/contrib/graphql/unified_trace_patcher.rbs

@@ -0,0 +1,11 @@
+module Datadog
+  module Tracing
+    module Contrib
+      module GraphQL
+        module UnifiedTracePatcher
+          def self?.patch!: (Array[GraphQL::Schema] schemas, Hash[Symbol, bool | Float | String | nil] options) -> void
+        end
+      end
+    end
+  end
+end

spec/datadog/tracing/contrib/graphql/configuration/settings_spec.rb

@@ -54,4 +54,30 @@
       end
     end
   end
+
+  describe 'with_unified_tracer' do
+    context 'when default' do
+      it do
+        settings = described_class.new
+
+        expect(settings.with_unified_tracer).to eq(false)
+      end
+    end
+
+    context 'when given `true`' do
+      it do
+        settings = described_class.new(with_unified_tracer: true)
+
+        expect(settings.with_unified_tracer).to eq(true)
+      end
+    end
+
+    context 'when given `false`' do
+      it do
+        settings = described_class.new(with_unified_tracer: false)
+
+        expect(settings.with_unified_tracer).to eq(false)
+      end
+    end
+  end
 end