Merge pull request #3672 from DataDog/vpellan/new-graphql-tracing

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

@@ -859,6 +859,7 @@ The `instrument :graphql` method accepts the following parameters. Additional op
 | ------------------------ | - | ------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------------- |
 | `enabled` | `DD_TRACE_GRAPHQL_ENABLED` | `Bool` | Whether the integration should create spans. | `true` |
 | `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`. Default is `false`, using `GraphQL::Tracing::DataDogTrace`                          | `false`          |
 | `service_name`           | | `String` | Service name used for graphql instrumentation                                                                                                                | `'ruby-graphql'` |
 
@@ -874,6 +875,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
@@ -886,6 +895,23 @@ end
 
 Do _NOT_ `instrument :graphql` in `Datadog.configure` if you choose to configure manually, as to avoid double tracing. These two means of configuring GraphQL tracing are considered mutually exclusive.
 
+**Adding custom tags to Datadog spans**
+
+You can add custom tags to Datadog spans by implementing the `prepare_span` method in a subclass, then manually configuring your schema.
+
+```ruby
+class YourSchema < GraphQL::Schema
+  module CustomTracing
+    include Datadog::Tracing::Contrib::GraphQL::UnifiedTrace
+    def prepare_span(trace_key, data, span)
+      span.set_tag("custom:#{trace_key}", data.keys.sort.join(","))
+    end
+  end
+  
+  trace_with CustomTracing
+end
+```
+
 ### gRPC
 
 The `grpc` integration adds both client and server interceptors, which run as middleware before executing the service's remote procedure call. As gRPC applications are often distributed, the integration shares trace information between client and server.

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,166 @@
+# 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.
+        # This tracer differs from the upstream one as it follows the unified naming convention specification,
+        # which is required to use features such as API Catalog.
+        # DEV-3.0: This tracer should be the default one in the next major version.
+        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.variables.instance_variable_get(:@storage).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_cache is initialized upstream, in ::GraphQL::Tracing::PlatformTrace
+            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)
+            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
+
+          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