SQL comment propagation for mysql2

lib/datadog/core/configuration/settings.rb

@@ -463,6 +463,16 @@ def initialize(*_)
             end
           end
 
+          option :sql_comment_propagation do |o|
+            o.default {
+              ENV.fetch(
+                Datadog::Tracing::Contrib::Propagation::SqlComment::Ext::ENV_SQL_COMMENT_PROPAGATION_MODE,
+                Datadog::Tracing::Contrib::Propagation::SqlComment::Ext::DISABLED
+              )
+            }
+            o.lazy
+          end
+
           # Enable trace collection and span generation.
           #
           # You can use this option to disable tracing without having to

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

@@ -4,6 +4,7 @@
 require_relative '../analytics'
 require_relative 'ext'
 require_relative '../ext'
+require_relative '../propagation/sql_comment'
 
 module Datadog
   module Tracing
@@ -39,6 +40,14 @@ def query(sql, options = {})
                 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])
+
+                propagation_mode = Contrib::Propagation::SqlComment::Mode.new(
+                  Datadog.configuration.tracing.sql_comment_propagation
+                )
+
+                Contrib::Propagation::SqlComment.annotate!(span, mode: propagation_mode)
+                sql = Contrib::Propagation::SqlComment.prepend_comment(sql, span, mode: propagation_mode)
+
                 super(sql, options)
               end
             end

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

@@ -0,0 +1,63 @@
+# typed: false
+
+require_relative 'sql_comment/comment'
+require_relative 'sql_comment/ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Propagation
+        module SqlComment
+          Mode = Struct.new(:mode) do
+            def enabled?
+              mode != Ext::DISABLED
+            end
+
+            def service?
+              mode == Ext::SERVICE
+            end
+
+            def full?
+              mode == Ext::FULL
+            end
+          end
+
+          def self.annotate!(span, mode:)
+            return unless mode.enabled?
+
+            span.set_tag('_dd.dbm_trace_injected', true) if mode.full?
+          end
+
+          def self.prepend_comment(sql, span, tags: {}, mode:)
+            return sql unless mode.enabled?
+
+            tags.merge!(service_context)     if mode.full? || mode.service?
+            tags.merge!(trace_context(span)) if mode.full?
+
+            "#{Comment.new(tags)} #{sql}"
+          end
+
+          private
+
+          def self.service_context
+            {
+              dde:  datadog_configuration.env,
+              ddps: datadog_configuration.service,
+              ddpv: datadog_configuration.version
+            }
+          end
+
+          def self.datadog_configuration
+            Datadog.configuration
+          end
+
+          def self.trace_context(_span)
+            {
+              traceparent: '00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01' # TODO: Derive from span?
+            }.freeze
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,37 @@
+# typed: false
+require 'erb'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Propagation
+        module SqlComment
+          class Comment
+            def initialize(hash)
+              @hash = hash
+            end
+
+            def to_s
+              @string ||= begin
+                ret = String.new
+
+                @hash.each do |key, value|
+                  next if value.nil?
+
+                  value = ERB::Util.url_encode(value) # url encode
+                  value.gsub!("'", "/'")              # escaping single quote
+
+                  ret << "#{key}='#{value}'," # escape sql
+                end
+
+                ret.chop!
+
+                "/*#{ret}*/"
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,19 @@
+# typed: false
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Propagation
+        module SqlComment
+          module Ext
+            ENV_SQL_COMMENT_PROPAGATION_MODE = 'DD_TRACE_SQL_COMMENT_PROPAGATION_MODE'.freeze
+
+            DISABLED = 'disabled'.freeze
+            SERVICE = 'service'.freeze
+            FULL = 'full'.freeze
+          end
+        end
+      end
+    end
+  end
+end

spec/datadog/tracing/contrib/mysql2/patcher_spec.rb

@@ -26,18 +26,23 @@
   let(:database) { ENV.fetch('TEST_MYSQL_DB') { 'mysql' } }
   let(:username) { ENV.fetch('TEST_MYSQL_USER') { 'root' } }
   let(:password) { ENV.fetch('TEST_MYSQL_PASSWORD') { 'root' } }
+  let(:sql_comment_propagation) { 'disabled' }
 
   before do
     Datadog.configure do |c|
+      c.service = 'my-service'
+      c.version = '2.0.0'
+      c.env = 'production'
+      c.tracing.sql_comment_propagation = sql_comment_propagation
       c.tracing.instrument :mysql2, configuration_options
     end
   end
 
   around do |example|
     # Reset before and after each example; don't allow global state to linger.
-    Datadog.registry[:mysql2].reset_configuration!
+    Datadog.configuration.reset!
     example.run
-    Datadog.registry[:mysql2].reset_configuration!
+    Datadog.configuration.reset!
   end
 
   describe 'tracing' do
@@ -103,6 +108,30 @@
             .to eq("Unknown column 'INVALID' in 'field list'")
         end
       end
+
+      context 'when sql comment propagation' do
+        context 'disabled' do
+          it do
+            client.query('SELECT 1')
+          end
+        end
+
+        context 'service' do
+          let(:sql_comment_propagation) { 'service' }
+
+          it do
+            client.query('SELECT 1')
+          end
+        end
+
+        context 'full' do
+          let(:sql_comment_propagation) { 'full' }
+
+          it do
+            client.query('SELECT 1')
+          end
+        end
+      end
     end
   end
 end