datadog 2.1.0 → 2.3.0

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'
+
+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,28 @@
+# frozen_string_literal: true
+
+if Gem.loaded_specs['graphql'] && Gem.loaded_specs['graphql'].version >= Gem::Version.new('2.0.19')
+  require_relative 'unified_trace'
+end
+
+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)
+            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/lograge/patcher.rb

@@ -20,6 +20,22 @@ module Datadog
 
           # patch applies our patch
           def patch
+            # ActiveSupport::TaggedLogging is the default Rails logger since Rails 5
+            if defined?(::ActiveSupport::TaggedLogging::Formatter) &&
+                ::Lograge::LogSubscribers::ActionController
+                    .logger&.formatter.is_a?(::ActiveSupport::TaggedLogging::Formatter)
+
+              Datadog.logger.error(
+                'Lograge and ActiveSupport::TaggedLogging (the default Rails log formatter) are not compatible: ' \
+                  'Lograge does not account for Rails log tags, creating polluted logs and breaking log formatting. ' \
+                  'Traces and Logs correlation may not work. ' \
+                  'Either: 1. Disable tagged logging in your Rails configuration ' \
+                  '`config.logger = ActiveSupport::Logger.new(STDOUT); ' \
+                  'config.active_job.logger = ActiveSupport::Logger.new(STDOUT)` ' \
+                  'or 2. Use the `semantic_logger` gem instead of `lograge`.'
+              )
+            end
+
             ::Lograge::LogSubscribers::Base.include(Instrumentation)
           end
         end

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

@@ -48,6 +48,11 @@ module Datadog
               o.default Contrib::Propagation::SqlComment::Ext::DISABLED
             end
 
+            option :append_comment do |o|
+              o.type :bool
+              o.default false
+            end
+
             option :peer_service do |o|
               o.type :string, nilable: true
               o.env Ext::ENV_PEER_SERVICE

data/lib/datadog/tracing/contrib/mysql2/instrumentation.rb

@@ -50,21 +50,14 @@ module Datadog
                 # Set analytics sample rate
                 Contrib::Analytics.set_sample_rate(span, analytics_sample_rate) if analytics_enabled?
 
+                span.set_tag(Contrib::Ext::DB::TAG_INSTANCE, query_options[:database])
                 span.set_tag(Ext::TAG_DB_NAME, query_options[:database])
                 span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_HOST, query_options[:host])
                 span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_PORT, query_options[:port])
 
                 Contrib::SpanAttributeSchema.set_peer_service!(span, Ext::PEER_SERVICE_SOURCES)
 
-                propagation_mode = Contrib::Propagation::SqlComment::Mode.new(comment_propagation)
-
-                Contrib::Propagation::SqlComment.annotate!(span, propagation_mode)
-                sql = Contrib::Propagation::SqlComment.prepend_comment(
-                  sql,
-                  span,
-                  trace_op,
-                  propagation_mode
-                )
+                sql = inject_propagation(span, sql, trace_op)
 
                 super(sql, options)
               end
@@ -72,6 +65,21 @@ module Datadog
 
             private
 
+            def inject_propagation(span, sql, trace_op)
+              propagation_mode = Contrib::Propagation::SqlComment::Mode.new(
+                datadog_configuration[:comment_propagation],
+                datadog_configuration[:append_comment]
+              )
+
+              Contrib::Propagation::SqlComment.annotate!(span, propagation_mode)
+              Contrib::Propagation::SqlComment.prepend_comment(
+                sql,
+                span,
+                trace_op,
+                propagation_mode
+              )
+            end
+
             def datadog_configuration
               Datadog.configuration.tracing[:mysql2]
             end
@@ -83,10 +91,6 @@ module Datadog
             def analytics_sample_rate
               datadog_configuration[:analytics_sample_rate]
             end
-
-            def comment_propagation
-              datadog_configuration[:comment_propagation]
-            end
           end
         end
       end

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

@@ -52,6 +52,11 @@ module Datadog
               o.default Contrib::Propagation::SqlComment::Ext::DISABLED
             end
 
+            option :append_comment do |o|
+              o.type :bool
+              o.default false
+            end
+
             option :peer_service do |o|
               o.type :string, nilable: true
               o.env Ext::ENV_PEER_SERVICE

data/lib/datadog/tracing/contrib/pg/instrumentation.rb

@@ -114,7 +114,10 @@ module Datadog
                 Contrib::Analytics.set_sample_rate(span, analytics_sample_rate) if analytics_enabled?
 
                 if sql
-                  propagation_mode = Contrib::Propagation::SqlComment::Mode.new(comment_propagation)
+                  propagation_mode = Contrib::Propagation::SqlComment::Mode.new(
+                    comment_propagation,
+                    datadog_configuration[:append_comment]
+                  )
                   Contrib::Propagation::SqlComment.annotate!(span, propagation_mode)
                   propagated_sql_statement = Contrib::Propagation::SqlComment.prepend_comment(
                     sql,

data/lib/datadog/tracing/contrib/propagation/sql_comment/ext.rb

@@ -20,10 +20,38 @@ module Datadog
             # The value should be `true` when `full` mode
             TAG_DBM_TRACE_INJECTED = '_dd.dbm_trace_injected'
 
+            # Database service/sql span service (i.e. the service executing the actual query)
+            #
+            # If fake services are disabled:
+            #   This value will be the same as the parent service
+            #
+            # If fake services are enabled:
+            #   This value is NOT the same as the parent service
+            #
+            # This should NOT be overridden by peer.service.
             KEY_DATABASE_SERVICE = 'dddbs'
+
+            # The global service environment (e.g. DD_ENV)
             KEY_ENVIRONMENT = 'dde'
+
+            # The global service name (e.g. DD_SERVICE)
             KEY_PARENT_SERVICE = 'ddps'
+
+            # The global service version (e.g. DD_VERSION)
             KEY_VERSION = 'ddpv'
+
+            # The hostname of the database server, as provided to the database client upon instantiation.
+            # @see Datadog::Tracing::Metadata::Ext::TAG_PEER_HOSTNAME
+            KEY_HOSTNAME = 'ddh'
+
+            # @see Datadog::Tracing::Contrib::Ext::DB::TAG_INSTANCE
+            KEY_DB_NAME = 'dddb'
+
+            # Users can use this attribute to specify the identity of the dependency/database they are connecting to.
+            # We should grab this attribute only if the user is EXPLICITLY specifying it.
+            # @see Datadog::Tracing::Metadata::Ext::TAG_PEER_SERVICE
+            KEY_PEER_SERVICE = 'ddprs'
+
             KEY_TRACEPARENT = 'traceparent'
           end
         end

data/lib/datadog/tracing/contrib/propagation/sql_comment/mode.rb

@@ -8,7 +8,7 @@ module Datadog
       module Propagation
         # Implements sql comment propagation related contracts.
         module SqlComment
-          Mode = Struct.new(:mode) do
+          Mode = Struct.new(:mode, :append) do
             def enabled?
               service? || full?
             end
@@ -20,6 +20,10 @@ module Datadog
             def full?
               mode == Ext::FULL
             end
+
+            def append?
+              append
+            end
           end
         end
       end

data/lib/datadog/tracing/contrib/propagation/sql_comment.rb

@@ -22,16 +22,28 @@ module Datadog
           def self.prepend_comment(sql, span_op, trace_op, mode)
             return sql unless mode.enabled?
 
+            config = Datadog.configuration
+
+            parent_service = config.service
+            peer_service = span_op.get_tag(Tracing::Metadata::Ext::TAG_PEER_SERVICE)
+
             tags = {
-              Ext::KEY_DATABASE_SERVICE => span_op.get_tag(Tracing::Metadata::Ext::TAG_PEER_SERVICE) || span_op.service,
-              Ext::KEY_ENVIRONMENT => datadog_configuration.env,
-              Ext::KEY_PARENT_SERVICE => datadog_configuration.service,
-              Ext::KEY_VERSION => datadog_configuration.version
+              Ext::KEY_ENVIRONMENT => config.env,
+              Ext::KEY_PARENT_SERVICE => parent_service,
+              Ext::KEY_VERSION => config.version,
+              Ext::KEY_HOSTNAME => span_op.get_tag(Tracing::Metadata::Ext::TAG_PEER_HOSTNAME),
+              Ext::KEY_DB_NAME => span_op.get_tag(Contrib::Ext::DB::TAG_INSTANCE),
+              Ext::KEY_PEER_SERVICE => peer_service,
             }
 
+            db_service = peer_service || span_op.service
+            if parent_service != db_service # Only set if it's different from parent_service; otherwise it's redundant
+              tags[Ext::KEY_DATABASE_SERVICE] = db_service
+            end
+
             if mode.full?
               # When tracing is disabled, trace_operation is a dummy object that does not contain data to build traceparent
-              if datadog_configuration.tracing.enabled
+              if config.tracing.enabled
                 tags[Ext::KEY_TRACEPARENT] =
                   Tracing::Distributed::TraceContext.new(fetcher: nil).send(:build_traceparent, trace_op.to_digest)
               else
@@ -42,11 +54,11 @@ module Datadog
               end
             end
 
-            "#{Comment.new(tags)} #{sql}"
-          end
-
-          def self.datadog_configuration
-            Datadog.configuration
+            if mode.append?
+              "#{sql} #{Comment.new(tags)}"
+            else
+              "#{Comment.new(tags)} #{sql}"
+            end
           end
         end
       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