feat: Add SQL Comment Propagator

.toys/.data/releases.yml

@@ -295,4 +295,4 @@ gems:
 
   - name: opentelemetry-sampler-xray
     directory: sampler/xray
-    version_constant: [OpenTelemetry, Sampler, XRay, VERSION]
+    version_constant: [OpenTelemetry, Sampler, XRay, VERSION]

helpers/sql-processor/Gemfile

@@ -11,6 +11,7 @@ gemspec
 group :test do
   gem 'bundler', '~> 2.4'
   gem 'minitest', '~> 5.0'
+  gem 'opentelemetry-sdk', '~> 1.5'
   gem 'opentelemetry-test-helpers', '~> 0.3'
   gem 'rake', '~> 13.0'
   gem 'rubocop', '~> 1.79.1'

helpers/sql-processor/lib/opentelemetry/helpers/sql_processor.rb

@@ -6,13 +6,15 @@
 
 require 'opentelemetry-common'
 require_relative 'sql_processor/obfuscator'
+require_relative 'sql_processor/commenter'
 
 module OpenTelemetry
   module Helpers
     # SQL processing utilities for OpenTelemetry instrumentation.
     #
     # This module provides a unified interface for SQL processing operations
-    # commonly needed in database adapter instrumentation, including SQL obfuscation.
+    # commonly needed in database adapter instrumentation, including SQL obfuscation
+    # and SQL comment-based trace context propagation.
     #
     # @api public
     module SqlProcessor

helpers/sql-processor/lib/opentelemetry/helpers/sql_processor/commenter.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require 'cgi'
+
+module OpenTelemetry
+  module Helpers
+    module SqlProcessor
+      # SqlCommenter provides SQL comment-based trace context propagation
+      # according to the SQL Commenter specification.
+      #
+      # This module implements a propagator interface compatible with Vitess,
+      # allowing it to be used as a drop-in replacement.
+      #
+      # @api public
+      module SqlCommenter
+        extend self
+
+        # SqlQuerySetter is responsible for formatting trace context as SQL comments
+        # and appending them to SQL queries according to the SQL Commenter specification.
+        #
+        # Format: /*key='value',key2='value2'*/
+        # Values are URL-encoded per the SQL Commenter spec
+        module SqlQuerySetter
+          extend self
+
+          # Appends trace context as a SQL comment to the carrier (SQL query string)
+          #
+          # @param carrier [String] The SQL query string to modify
+          # @param headers [Hash] Hash of trace context headers (e.g., {'traceparent' => '00-...'})
+          def set(carrier, headers)
+            return if headers.empty?
+            return if carrier.frozen?
+
+            # Convert headers hash to SQL commenter format
+            # Format: /*key1='value1',key2='value2'*/
+            comment_parts = headers.map do |key, value|
+              # URL encode values as per SQL Commenter spec (using URI component encoding)
+              encoded_value = CGI.escapeURIComponent(value.to_s)
+              "#{key}='#{encoded_value}'"
+            end
+
+            # Append to end of query (spec recommendation)
+            carrier << " /*#{comment_parts.join(',')}*/"
+          end
+        end
+
+        # SqlQueryPropagator propagates trace context using SQL comments
+        # according to the SQL Commenter specification.
+        #
+        # This propagator implements the same interface as the Vitess propagator
+        # and can be used as a drop-in replacement.
+        #
+        # @example
+        #   propagator = OpenTelemetry::Helpers::SqlProcessor::SqlCommenter.sql_query_propagator
+        #   sql = "SELECT * FROM users"
+        #   propagator.inject(sql, context: current_context)
+        #   # => "SELECT * FROM users /*traceparent='00-...',tracestate='...'*/"
+        module SqlQueryPropagator
+          extend self
+
+          # Injects trace context into a SQL query as a comment
+          #
+          # @param carrier [String] The SQL query string to inject context into
+          # @param context [optional, Context] The context to inject. Defaults to current context.
+          # @param setter [optional, #set] The setter to use for appending the comment.
+          #   Defaults to SqlQuerySetter.
+          # @return [nil]
+          def inject(carrier, context: OpenTelemetry::Context.current, setter: SqlQuerySetter)
+            # Use the global propagator to extract headers into a hash
+            headers = {}
+            OpenTelemetry.propagation.inject(headers, context: context)
+
+            # Pass the headers to our SQL comment setter
+            setter.set(carrier, headers)
+            nil
+          end
+        end
+
+        # Returns the SqlQueryPropagator module for stateless propagation
+        #
+        # @return [Module] The SqlQueryPropagator module
+        def sql_query_propagator
+          SqlQueryPropagator
+        end
+      end
+    end
+  end
+end

helpers/sql-processor/opentelemetry-helpers-sql-processor.gemspec

@@ -25,6 +25,7 @@ Gem::Specification.new do |spec|
   spec.require_paths = ['lib']
   spec.required_ruby_version = '>= 3.2'
 
+  spec.add_dependency 'opentelemetry-api', '~> 1.0'
   spec.add_dependency 'opentelemetry-common', '~> 0.21'
 
   if spec.respond_to?(:metadata)

helpers/sql-processor/test/opentelemetry/helpers/sql_processor/commenter_test.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require 'test_helper'
+
+describe OpenTelemetry::Helpers::SqlProcessor::SqlCommenter do
+  let(:span_id) { 'e457b5a2e4d86bd1' }
+  let(:trace_id) { '80f198ee56343ba864fe8b2a57d3eff7' }
+  let(:trace_flags) { OpenTelemetry::Trace::TraceFlags::SAMPLED }
+
+  let(:context) do
+    OpenTelemetry::Trace.context_with_span(
+      OpenTelemetry::Trace.non_recording_span(
+        OpenTelemetry::Trace::SpanContext.new(
+          trace_id: Array(trace_id).pack('H*'),
+          span_id: Array(span_id).pack('H*'),
+          trace_flags: trace_flags
+        )
+      )
+    )
+  end
+
+  describe 'SqlQueryPropagator.inject' do
+    let(:propagator) { OpenTelemetry::Helpers::SqlProcessor::SqlCommenter.sql_query_propagator }
+
+    it 'injects trace context into SQL' do
+      sql = +'SELECT * FROM users'
+      propagator.inject(sql, context: context)
+
+      expected = "SELECT * FROM users /*traceparent='00-#{trace_id}-#{span_id}-01'*/"
+      _(sql).must_equal(expected)
+    end
+
+    it 'handles frozen strings by not modifying them' do
+      sql = -'SELECT * FROM users'
+      propagator.inject(sql, context: context)
+
+      # Frozen string should remain unchanged (setter will return early)
+      _(sql).must_equal('SELECT * FROM users')
+    end
+
+    it 'handles empty context' do
+      sql = +'SELECT * FROM users'
+      propagator.inject(sql, context: OpenTelemetry::Context.empty)
+
+      # Should not modify SQL when context produces no headers
+      _(sql).must_equal('SELECT * FROM users')
+    end
+
+    it 'includes tracestate when present' do
+      span_context = OpenTelemetry::Trace::SpanContext.new(
+        trace_id: Array(trace_id).pack('H*'),
+        span_id: Array(span_id).pack('H*'),
+        trace_flags: trace_flags,
+        tracestate: OpenTelemetry::Trace::Tracestate.from_string('congo=t61rcWkgMzE,rojo=00f067aa0ba902b7')
+      )
+
+      ctx = OpenTelemetry::Trace.context_with_span(
+        OpenTelemetry::Trace.non_recording_span(span_context)
+      )
+
+      sql = +'SELECT * FROM users'
+      propagator.inject(sql, context: ctx)
+
+      expected = "SELECT * FROM users /*traceparent='00-#{trace_id}-#{span_id}-01',tracestate='congo%3Dt61rcWkgMzE%2Crojo%3D00f067aa0ba902b7'*/"
+      _(sql).must_equal(expected)
+    end
+
+    it 'returns nil' do
+      sql = +'SELECT * FROM users'
+      result = propagator.inject(sql, context: context)
+
+      _(result).must_be_nil
+    end
+  end
+
+  describe 'SqlQuerySetter.set' do
+    let(:setter) { OpenTelemetry::Helpers::SqlProcessor::SqlCommenter::SqlQuerySetter }
+
+    it 'formats headers as SQL comments' do
+      sql = +'SELECT * FROM users'
+      headers = { 'traceparent' => '00-80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-01' }
+
+      setter.set(sql, headers)
+
+      expected = "SELECT * FROM users /*traceparent='00-80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-01'*/"
+      _(sql).must_equal(expected)
+    end
+
+    it 'URL encodes values' do
+      sql = +'SELECT * FROM users'
+      headers = { 'key' => 'value with spaces' }
+
+      setter.set(sql, headers)
+
+      expected = "SELECT * FROM users /*key='value%20with%20spaces'*/"
+      _(sql).must_equal(expected)
+    end
+
+    it 'handles empty headers' do
+      sql = +'SELECT * FROM users'
+      setter.set(sql, {})
+
+      _(sql).must_equal('SELECT * FROM users')
+    end
+
+    it 'handles frozen strings by not modifying them' do
+      sql = -'SELECT * FROM users'
+      headers = { 'traceparent' => '00-80f198ee56343ba864fe8b2a57d3eff7-e457b5a2e4d86bd1-01' }
+
+      setter.set(sql, headers)
+
+      # Frozen string should remain unchanged
+      _(sql).must_equal('SELECT * FROM users')
+    end
+  end
+end

helpers/sql-processor/test/test_helper.rb

@@ -10,3 +10,9 @@
 
 require 'minitest/autorun'
 require 'opentelemetry-helpers-sql-processor'
+require 'opentelemetry/sdk'
+
+OpenTelemetry.logger = Logger.new($stderr, level: ENV.fetch('OTEL_LOG_LEVEL', 'fatal').to_sym)
+
+# Configure the SDK to set up the default propagators
+OpenTelemetry::SDK.configure

instrumentation/mysql2/lib/opentelemetry/instrumentation/mysql2/instrumentation.rb

@@ -10,9 +10,10 @@ module Mysql2
       # The Instrumentation class contains logic to detect and install the Mysql2
       # instrumentation
       class Instrumentation < OpenTelemetry::Instrumentation::Base
-        install do |_config|
+        install do |config|
           require_dependencies
           patch_client
+          configure_propagator(config)
         end
 
         present do
@@ -23,6 +24,9 @@ class Instrumentation < OpenTelemetry::Instrumentation::Base
         option :db_statement, default: :obfuscate, validate: %I[omit include obfuscate]
         option :span_name, default: :statement_type, validate: %I[statement_type db_name db_operation_and_name]
         option :obfuscation_limit, default: 2000, validate: :integer
+        option :propagator, default: 'none', validate: %w[none tracecontext vitess]
+
+        attr_reader :propagator
 
         private
 
@@ -33,6 +37,24 @@ def require_dependencies
         def patch_client
           ::Mysql2::Client.prepend(Patches::Client)
         end
+
+        def configure_propagator(config)
+          propagator = config[:propagator]
+          @propagator = case propagator
+                        when 'tracecontext' then OpenTelemetry::Helpers::SqlProcessor::SqlCommenter.sql_query_propagator
+                        when 'vitess' then fetch_propagator(propagator, 'OpenTelemetry::Propagator::Vitess')
+                        when 'none', nil then nil
+                        else
+                          OpenTelemetry.logger.warn "The #{propagator} propagator is unknown and cannot be configured"
+                        end
+        end
+
+        def fetch_propagator(name, class_name, gem_suffix = name)
+          Kernel.const_get(class_name).sql_query_propagator
+        rescue NameError
+          OpenTelemetry.logger.warn "The #{name} propagator cannot be configured - please add opentelemetry-helpers-#{gem_suffix} to your Gemfile"
+          nil
+        end
       end
     end
   end

instrumentation/mysql2/lib/opentelemetry/instrumentation/mysql2/patches/client.rb

@@ -18,8 +18,16 @@ def query(sql, options = {})
               _otel_span_name(sql),
               attributes: _otel_span_attributes(sql),
               kind: :client
-            ) do
-              super
+            ) do |_span, context|
+              if propagator && sql.frozen?
+                sql = +sql
+                propagator.inject(sql, context: context)
+                sql.freeze
+              elsif propagator
+                propagator.inject(sql, context: context)
+              end
+
+              super(sql, options)
             end
           end
 
@@ -28,8 +36,16 @@ def prepare(sql)
               _otel_span_name(sql),
               attributes: _otel_span_attributes(sql),
               kind: :client
-            ) do
-              super
+            ) do |_span, context|
+              if propagator && sql.frozen?
+                sql = +sql
+                propagator.inject(sql, context: context)
+                sql.freeze
+              elsif propagator
+                propagator.inject(sql, context: context)
+              end
+
+              super(sql)
             end
           end
 
@@ -93,6 +109,10 @@ def tracer
           def config
             Mysql2::Instrumentation.instance.config
           end
+
+          def propagator
+            Mysql2::Instrumentation.instance.propagator
+          end
         end
       end
     end