datadog 2.0.0 → 2.2.0

data/lib/datadog/tracing/contrib/active_support/notifications/event.rb

@@ -26,21 +26,25 @@ module Datadog
                 super
               end
 
-              def subscription(span_name = nil, options = nil)
+              def subscription(span_name = nil, span_options = nil, on_start: nil, on_finish: nil, trace: nil)
                 super(
                   span_name || self.span_name,
-                  options || span_options,
-                  &method(:process)
+                  span_options || self.span_options,
+                  on_start: on_start,
+                  on_finish: on_finish,
+                  trace: trace
                 )
               end
 
-              def subscribe(pattern = nil, span_name = nil, options = nil)
+              def subscribe(pattern = nil, span_name = nil, span_options = nil)
                 if supported?
                   super(
                     pattern || event_name,
                     span_name || self.span_name,
-                    options || span_options,
-                    &method(:process)
+                    span_options || self.span_options,
+                    on_start: method(:on_start),
+                    on_finish: method(:on_finish),
+                    trace: method(:trace?)
                   )
                 end
               end
@@ -62,6 +66,25 @@ module Datadog
                 payload[:exception_object] ||
                   payload[:exception] # Fallback for ActiveSupport < 5.0
               end
+
+              def on_start(_span, _event, _id, _payload); end
+
+              def on_finish(span, _event, _id, payload)
+                record_exception(span, payload)
+              end
+
+              def trace?(_event, _payload)
+                true
+              end
+
+              def record_exception(span, payload)
+                if payload[:exception_object]
+                  span.set_error(payload[:exception_object])
+                elsif payload[:exception]
+                  # Fallback for ActiveSupport < 5.0
+                  span.set_error(payload[:exception])
+                end
+              end
             end
           end
         end

data/lib/datadog/tracing/contrib/active_support/notifications/subscriber.rb

@@ -45,16 +45,28 @@ module Datadog
               end
 
               # Creates a subscription and immediately activates it.
-              def subscribe(pattern, span_name, options = {}, &block)
-                subscription(span_name, options, &block).tap do |subscription|
+              def subscribe(pattern, span_name, span_options = {}, on_start: nil, on_finish: nil, trace: nil)
+                subscription(
+                  span_name,
+                  span_options,
+                  on_start: on_start,
+                  on_finish: on_finish,
+                  trace: trace
+                ).tap do |subscription|
                   subscription.subscribe(pattern)
                 end
               end
 
               # Creates a subscription without activating it.
               # Subscription is added to the inheriting class' list of subscriptions.
-              def subscription(span_name, options = {}, &block)
-                Subscription.new(span_name, options, &block).tap do |subscription|
+              def subscription(span_name, span_options = {}, on_start: nil, on_finish: nil, trace: nil)
+                Subscription.new(
+                  span_name,
+                  span_options,
+                  on_start: on_start,
+                  on_finish: on_finish,
+                  trace: trace
+                ).tap do |subscription|
                   subscriptions << subscription
                 end
               end

data/lib/datadog/tracing/contrib/active_support/notifications/subscription.rb

@@ -9,31 +9,34 @@ module Datadog
           class Subscription
             attr_accessor \
               :span_name,
-              :options
-
-            def initialize(span_name, options, &block)
-              raise ArgumentError, 'Must be given a block!' unless block
+              :span_options
+
+            # @param span_name [String] the operation name for the span
+            # @param span_options [Hash] span_options to pass during span creation
+            # @param on_start [Proc] a block to run when the event is fired,
+            #   might not include all required information in the `payload` argument.
+            # @param on_finish [Proc] a block to run when the event has finished processing,
+            #   possibly including more information in the `payload` argument.
+            # @param trace [Proc] whether to trace the event. Defaults to returning `true`.
+            def initialize(span_name, span_options, on_start: nil, on_finish: nil, trace: nil)
+              raise ArgumentError, 'Must be given either on_start or on_finish' unless on_start || on_finish
 
               @span_name = span_name
-              @options = options
-              @handler = Handler.new(&block)
+              @span_options = span_options
+              @on_start = Handler.new(on_start)
+              @on_finish = Handler.new(on_finish)
+              @trace = trace
               @callbacks = Callbacks.new
             end
 
-            # ActiveSupport 3.x calls this
-            def call(name, start, finish, id, payload)
-              start_span(name, id, payload, start)
-              finish_span(name, id, payload, finish)
-            end
-
-            # ActiveSupport 4+ calls this on start
+            # Called by ActiveSupport on event start
             def start(name, id, payload)
-              start_span(name, id, payload)
+              start_span(name, id, payload) if @trace&.call(name, payload)
             end
 
-            # ActiveSupport 4+ calls this on finish
+            # Called by ActiveSupport on event finish
             def finish(name, id, payload)
-              finish_span(name, id, payload)
+              finish_span(name, id, payload) if payload[:datadog_span]
             end
 
             def before_trace(&block)
@@ -69,7 +72,8 @@ module Datadog
             protected
 
             attr_reader \
-              :handler,
+              :on_start,
+              :on_finish,
               :callbacks
 
             def start_span(name, id, payload, start = nil)
@@ -77,11 +81,15 @@ module Datadog
               callbacks.run(name, :before_trace, id, payload, start)
 
               # Start a trace
-              Tracing.trace(@span_name, **@options).tap do |span|
-                # Start span if time is provided
-                span.start(start) unless start.nil?
-                payload[:datadog_span] = span
-              end
+              span = Tracing.trace(@span_name, **@span_options)
+
+              # Start span if time is provided
+              span.start(start) unless start.nil?
+              payload[:datadog_span] = span
+
+              on_start.run(span, name, id, payload)
+
+              span
             end
 
             def finish_span(name, id, payload, finish = nil)
@@ -90,7 +98,7 @@ module Datadog
                 return nil if span.nil?
 
                 # Run handler for event
-                handler.run(span, name, id, payload)
+                on_finish.run(span, name, id, payload)
 
                 # Finish the span
                 span.finish(finish)
@@ -109,21 +117,17 @@ module Datadog
             class Handler
               attr_reader :block
 
-              def initialize(&block)
+              def initialize(block)
                 @block = block
               end
 
               def run(span, name, id, payload)
-                run!(span, name, id, payload)
+                @block.call(span, name, id, payload) if @block
               rescue StandardError => e
                 Datadog.logger.debug(
                   "ActiveSupport::Notifications handler for '#{name}' failed: #{e.class.name} #{e.message}"
                 )
               end
-
-              def run!(*args)
-                @block.call(*args)
-              end
             end
 
             # Wrapper for subscription callbacks

data/lib/datadog/tracing/contrib/analytics.rb

@@ -9,6 +9,11 @@ module Datadog
       module Analytics
         module_function
 
+        # Applies Analytics sampling rate, if applicable for this Contrib::Configuration.
+        def set_rate!(span, configuration)
+          set_sample_rate(span, configuration[:analytics_sample_rate]) if enabled?(configuration[:analytics_enabled])
+        end
+
         # Checks whether analytics should be enabled.
         # `flag` is a truthy/falsey value that represents a setting on the integration.
         def enabled?(flag = nil)

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

@@ -42,6 +42,11 @@ module Datadog
               o.type :bool
               o.default false
             end
+
+            option :with_unified_tracer do |o|
+              o.type :bool
+              o.default false
+            end
           end
         end
       end

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

@@ -4,6 +4,7 @@ require_relative '../analytics'
 require_relative '../patcher'
 require_relative 'tracing_patcher'
 require_relative 'trace_patcher'
+require_relative 'unified_trace_patcher'
 
 module Datadog
   module Tracing
@@ -23,10 +24,15 @@ module Datadog
             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)

data/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

data/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

data/lib/datadog/tracing/contrib/kafka/consumer_event.rb

@@ -6,7 +6,7 @@ module Datadog
       module Kafka
         # Defines basic behaviors for an event for a consumer.
         module ConsumerEvent
-          def process(span, _event, _id, payload)
+          def on_start(span, _event, _id, payload)
             super
 
             span.set_tag(Ext::TAG_GROUP, payload[:group_id])

data/lib/datadog/tracing/contrib/kafka/consumer_group_event.rb

@@ -6,7 +6,7 @@ module Datadog
       module Kafka
         # Defines basic behaviors for an event for a consumer group.
         module ConsumerGroupEvent
-          def process(span, _event, _id, payload)
+          def on_start(span, _event, _id, payload)
             super
 
             span.resource = payload[:group_id]

data/lib/datadog/tracing/contrib/kafka/event.rb

@@ -29,7 +29,7 @@ module Datadog
               Datadog.configuration.tracing[:kafka]
             end
 
-            def process(span, _event, _id, payload)
+            def on_start(span, _event, _id, payload)
               span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
               span.set_tag(Contrib::Ext::Messaging::TAG_SYSTEM, Ext::TAG_MESSAGING_SYSTEM)
 

data/lib/datadog/tracing/contrib/kafka/events/connection/request.rb

@@ -15,7 +15,9 @@ module Datadog
 
               EVENT_NAME = 'request.connection.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.resource = payload[:api]
@@ -24,8 +26,6 @@ module Datadog
                 span.set_tag(Ext::TAG_RESPONSE_SIZE, payload[:response_size]) if payload.key?(:response_size)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_CONNECTION_REQUEST
               end

data/lib/datadog/tracing/contrib/kafka/events/consumer/process_batch.rb

@@ -17,7 +17,9 @@ module Datadog
 
               EVENT_NAME = 'process_batch.consumer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.resource = payload[:topic]
@@ -31,8 +33,6 @@ module Datadog
                 span.set_tag(Ext::TAG_OFFSET_LAG, payload[:offset_lag]) if payload.key?(:offset_lag)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_PROCESS_BATCH
               end

data/lib/datadog/tracing/contrib/kafka/events/consumer/process_message.rb

@@ -17,7 +17,9 @@ module Datadog
 
               EVENT_NAME = 'process_message.consumer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.resource = payload[:topic]
@@ -29,8 +31,6 @@ module Datadog
                 span.set_tag(Ext::TAG_OFFSET_LAG, payload[:offset_lag]) if payload.key?(:offset_lag)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_PROCESS_MESSAGE
               end

data/lib/datadog/tracing/contrib/kafka/events/consumer_group/heartbeat.rb

@@ -19,7 +19,9 @@ module Datadog
 
               EVENT_NAME = 'heartbeat.consumer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 if payload.key?(:topic_partitions)
@@ -29,8 +31,6 @@ module Datadog
                 end
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_CONSUMER_HEARTBEAT
               end

data/lib/datadog/tracing/contrib/kafka/events/produce_operation/send_messages.rb

@@ -15,7 +15,9 @@ module Datadog
 
               EVENT_NAME = 'send_messages.producer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.set_tag(Ext::TAG_MESSAGE_COUNT, payload[:message_count]) if payload.key?(:message_count)
@@ -23,8 +25,6 @@ module Datadog
                 span.set_tag(Tracing::Metadata::Ext::TAG_KIND, Tracing::Metadata::Ext::SpanKind::TAG_PRODUCER)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_SEND_MESSAGES
               end

data/lib/datadog/tracing/contrib/kafka/events/producer/deliver_messages.rb

@@ -15,7 +15,9 @@ module Datadog
 
               EVENT_NAME = 'deliver_messages.producer.kafka'
 
-              def self.process(span, _event, _id, payload)
+              module_function
+
+              def on_start(span, _event, _id, payload)
                 super
 
                 span.set_tag(Ext::TAG_ATTEMPTS, payload[:attempts]) if payload.key?(:attempts)
@@ -26,8 +28,6 @@ module Datadog
                 span.set_tag(Tracing::Metadata::Ext::TAG_KIND, Tracing::Metadata::Ext::SpanKind::TAG_PRODUCER)
               end
 
-              module_function
-
               def span_name
                 Ext::SPAN_DELIVER_MESSAGES
               end

data/lib/datadog/tracing/contrib/racecar/event.rb

@@ -19,7 +19,7 @@ module Datadog
           # Class methods for Racecar events.
           # Note, they share the same process method and before_trace method.
           module ClassMethods
-            def subscription(*args)
+            def subscription(*args, **kwargs)
               super.tap do |subscription|
                 subscription.before_trace { ensure_clean_context! }
               end
@@ -33,7 +33,7 @@ module Datadog
               Datadog.configuration.tracing[:racecar]
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               span.service = configuration[:service_name]
               span.resource = payload[:consumer_class]
 

data/lib/datadog/tracing/contrib/rails/ext.rb

@@ -14,6 +14,15 @@ module Datadog
           ENV_ANALYTICS_SAMPLE_RATE = 'DD_TRACE_RAILS_ANALYTICS_SAMPLE_RATE'
           ENV_DISABLE = 'DISABLE_DATADOG_RAILS'
 
+          SPAN_RUNNER_FILE = 'rails.runner.file'
+          SPAN_RUNNER_INLINE = 'rails.runner.inline'
+          SPAN_RUNNER_STDIN = 'rails.runner.stdin'
+          TAG_COMPONENT = 'rails'
+          TAG_OPERATION_FILE = 'runner.file'
+          TAG_OPERATION_INLINE = 'runner.inline'
+          TAG_OPERATION_STDIN = 'runner.stdin'
+          TAG_RUNNER_SOURCE = 'source'
+
           # @!visibility private
           MINIMUM_VERSION = Gem::Version.new('4')
         end

data/lib/datadog/tracing/contrib/rails/patcher.rb

@@ -5,6 +5,7 @@ require_relative '../rack/middlewares'
 require_relative 'framework'
 require_relative 'log_injection'
 require_relative 'middlewares'
+require_relative 'runner'
 require_relative 'utils'
 require_relative '../semantic_logger/patcher'
 
@@ -28,6 +29,7 @@ module Datadog
           def patch
             patch_before_initialize
             patch_after_initialize
+            patch_rails_runner
           end
 
           def patch_before_initialize
@@ -81,6 +83,11 @@ module Datadog
           def setup_tracer
             Contrib::Rails::Framework.setup
           end
+
+          # Instruments the `bin/rails runner` command.
+          def patch_rails_runner
+            ::Rails::Command.singleton_class.prepend(Command) if defined?(::Rails::Command)
+          end
         end
       end
     end