Move append to Mode

DataDog · Jul 26, 2024 · 26b1b32 · 26b1b32
1 parent de31dbb
commit 26b1b32
Show file tree

Hide file tree

Showing 8 changed files with 33 additions and 27 deletions.
diff --git a/lib/datadog/tracing/contrib/mysql2/instrumentation.rb b/lib/datadog/tracing/contrib/mysql2/instrumentation.rb
@@ -66,15 +66,14 @@ def query(sql, options = {})
             private
 
             def inject_propagation(span, sql, trace_op)
-              propagation_mode = Contrib::Propagation::SqlComment::Mode.new(datadog_configuration[:comment_propagation])
+              propagation_mode = Contrib::Propagation::SqlComment::Mode.new(datadog_configuration[:comment_propagation], datadog_configuration[:append])
 
               Contrib::Propagation::SqlComment.annotate!(span, propagation_mode)
               Contrib::Propagation::SqlComment.prepend_comment(
                 sql,
                 span,
                 trace_op,
-                propagation_mode,
-                datadog_configuration[:append]
+                propagation_mode
               )
             end
 

diff --git a/lib/datadog/tracing/contrib/pg/instrumentation.rb b/lib/datadog/tracing/contrib/pg/instrumentation.rb
@@ -114,14 +114,13 @@ def trace(name, sql: nil, statement_name: nil, block: nil)
                 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])
                   Contrib::Propagation::SqlComment.annotate!(span, propagation_mode)
                   propagated_sql_statement = Contrib::Propagation::SqlComment.prepend_comment(
                     sql,
                     span,
                     trace_op,
-                    propagation_mode,
-                    datadog_configuration[:append]
+                    propagation_mode
                   )
                 end
 

diff --git a/lib/datadog/tracing/contrib/propagation/sql_comment.rb b/lib/datadog/tracing/contrib/propagation/sql_comment.rb
@@ -19,7 +19,7 @@ def self.annotate!(span_op, mode)
 
           # Inject span_op and trace_op instead of TraceDigest to improve memory usage
           # for `disabled` and `service` mode
-          def self.prepend_comment(sql, span_op, trace_op, mode, append)
+          def self.prepend_comment(sql, span_op, trace_op, mode)
             return sql unless mode.enabled?
 
             config = Datadog.configuration
@@ -54,7 +54,7 @@ def self.prepend_comment(sql, span_op, trace_op, mode, append)
               end
             end
 
-            if append
+            if mode.append?
               "#{sql} #{Comment.new(tags)}"
             else
               "#{Comment.new(tags)} #{sql}"

diff --git a/lib/datadog/tracing/contrib/propagation/sql_comment/mode.rb b/lib/datadog/tracing/contrib/propagation/sql_comment/mode.rb
@@ -8,17 +8,25 @@ module Contrib
       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
 
             def service?
-              mode == Ext::SERVICE
+              @mode == Ext::SERVICE
             end
 
             def full?
-              mode == Ext::FULL
+              @mode == Ext::FULL
+            end
+
+            def prepend?
+              !@append
+            end
+
+            def append?
+              @append
             end
           end
         end

diff --git a/lib/datadog/tracing/contrib/trilogy/instrumentation.rb b/lib/datadog/tracing/contrib/trilogy/instrumentation.rb
@@ -56,15 +56,14 @@ def query(sql)
 
                 Contrib::SpanAttributeSchema.set_peer_service!(span, Ext::PEER_SERVICE_SOURCES)
 
-                propagation_mode = Contrib::Propagation::SqlComment::Mode.new(comment_propagation)
+                propagation_mode = Contrib::Propagation::SqlComment::Mode.new(comment_propagation, datadog_configuration[:append])
 
                 Contrib::Propagation::SqlComment.annotate!(span, propagation_mode)
                 sql = Contrib::Propagation::SqlComment.prepend_comment(
                   sql,
                   span,
                   trace_op,
-                  propagation_mode,
-                  datadog_configuration[:append]
+                  propagation_mode
                 )
 
                 super(sql)

diff --git a/sig/datadog/tracing/contrib/propagation/sql_comment/mode.rbs b/sig/datadog/tracing/contrib/propagation/sql_comment/mode.rbs
@@ -3,12 +3,14 @@ module Datadog
     module Contrib
       module Propagation
         module SqlComment
-          class Mode < ::Struct[String]
-            attr_accessor mode: String
+          class Mode < ::Struct[untyped]
+            def initialize: (string mode, bool append) -> void
 
             def enabled?: -> bool
             def service?: -> bool
             def full?: -> bool
+            def prepend?: -> bool
+            def append?: -> bool
           end
         end
       end

diff --git a/spec/datadog/tracing/contrib/propagation/sql_comment_spec.rb b/spec/datadog/tracing/contrib/propagation/sql_comment_spec.rb
@@ -2,7 +2,8 @@
 require 'datadog/tracing/contrib/propagation/sql_comment/mode'
 
 RSpec.describe Datadog::Tracing::Contrib::Propagation::SqlComment do
-  let(:propagation_mode) { Datadog::Tracing::Contrib::Propagation::SqlComment::Mode.new(mode) }
+  let(:propagation_mode) { Datadog::Tracing::Contrib::Propagation::SqlComment::Mode.new(mode, append) }
+  let(:append) { false }
 
   describe '.annotate!' do
     let(:span_op) { Datadog::Tracing::SpanOperation.new('sql_comment_propagation_span', service: 'db_service') }
@@ -73,7 +74,7 @@
         )
       end
 
-      subject { described_class.prepend_comment(sql_statement, span_op, trace_op, propagation_mode, append) }
+      subject { described_class.prepend_comment(sql_statement, span_op, trace_op, propagation_mode) }
 
       context 'when `disabled` mode' do
         let(:mode) { 'disabled' }
@@ -222,7 +223,7 @@
         Datadog::Tracing.trace('dummy.sql') do |span_op, trace_op|
           span_op.service = 'db_service'
 
-          result = described_class.prepend_comment(sql_statement, span_op, trace_op, propagation_mode, append)
+          result = described_class.prepend_comment(sql_statement, span_op, trace_op, propagation_mode)
         end
 
         result

diff --git a/spec/datadog/tracing/contrib/sql_comment_propagation_examples.rb b/spec/datadog/tracing/contrib/sql_comment_propagation_examples.rb
@@ -1,7 +1,7 @@
 RSpec.shared_examples_for 'with sql comment propagation' do |span_op_name:, error: nil|
   context 'when default `disabled`' do
     it_behaves_like 'propagates with sql comment', mode: 'disabled', span_op_name: span_op_name, error: error do
-      let(:propagation_mode) { Datadog::Tracing::Contrib::Propagation::SqlComment::Mode.new('disabled') }
+      let(:propagation_mode) { Datadog::Tracing::Contrib::Propagation::SqlComment::Mode.new('disabled', append) }
     end
   end
 
@@ -14,7 +14,7 @@
     end
 
     it_behaves_like 'propagates with sql comment', mode: 'service', span_op_name: span_op_name, error: error do
-      let(:propagation_mode) { Datadog::Tracing::Contrib::Propagation::SqlComment::Mode.new('service') }
+      let(:propagation_mode) { Datadog::Tracing::Contrib::Propagation::SqlComment::Mode.new('service', append) }
     end
   end
 
@@ -25,7 +25,7 @@
       end
 
       it_behaves_like 'propagates with sql comment', mode: mode, span_op_name: span_op_name, error: error do
-        let(:propagation_mode) { Datadog::Tracing::Contrib::Propagation::SqlComment::Mode.new(mode) }
+        let(:propagation_mode) { Datadog::Tracing::Contrib::Propagation::SqlComment::Mode.new(mode, append) }
       end
     end
   end
@@ -36,7 +36,7 @@
 
   it "propagates with mode: #{mode}" do
     expect(Datadog::Tracing::Contrib::Propagation::SqlComment::Mode)
-      .to receive(:new).with(mode).and_return(propagation_mode)
+      .to receive(:new).with(mode, append).and_return(propagation_mode)
 
     if error
       expect { subject }.to raise_error(error)
@@ -70,8 +70,7 @@
       sql_statement,
       a_span_operation_with(service: service_name),
       duck_type(:to_digest),
-      propagation_mode,
-      append
+      propagation_mode
     )
   end
 
@@ -92,8 +91,7 @@
         sql_statement,
         a_span_operation_with(service: service_name),
         duck_type(:to_digest),
-        propagation_mode,
-        append
+        propagation_mode
       )
     end
   end