From a2669ad5d86854052d3c9a884fdcc37ab0fbc50c Mon Sep 17 00:00:00 2001
From: Sam <370182+plantfansam@users.noreply.github.com>
Date: Tue, 20 Feb 2024 09:23:07 -0700
Subject: [PATCH] fix: dup string if frozen in trilogy query (#863)

* fix: propagate context from trilogy instrumentation for frozen strings
---
 .../trilogy/instrumentation.rb                | 19 +-------
 .../instrumentation/trilogy/patches/client.rb |  9 +++-
 .../trilogy/instrumentation_test.rb           | 44 +++++++++++++++++--
 3 files changed, 50 insertions(+), 22 deletions(-)

diff --git a/instrumentation/trilogy/lib/opentelemetry/instrumentation/trilogy/instrumentation.rb b/instrumentation/trilogy/lib/opentelemetry/instrumentation/trilogy/instrumentation.rb
index cdb411b96..684fbe660 100644
--- a/instrumentation/trilogy/lib/opentelemetry/instrumentation/trilogy/instrumentation.rb
+++ b/instrumentation/trilogy/lib/opentelemetry/instrumentation/trilogy/instrumentation.rb
@@ -7,22 +7,6 @@
 module OpenTelemetry
   module Instrumentation
     module Trilogy
-      # @api private
-      class NoopPropagator
-        EMPTY_LIST = [].freeze
-        private_constant(:EMPTY_LIST)
-
-        def inject(carrier, context: Context.current, setter: nil); end
-
-        def extract(carrier, context: Context.current, getter: nil)
-          context
-        end
-
-        def fields
-          EMPTY_LIST
-        end
-      end
-
       # The Instrumentation class contains logic to detect and install the Trilogy instrumentation
       class Instrumentation < OpenTelemetry::Instrumentation::Base
         install do |config|
@@ -61,10 +45,9 @@ def configure_propagator(config)
           propagator = config[:propagator]
           @propagator = case propagator
                         when 'vitess' then fetch_propagator(propagator, 'OpenTelemetry::Propagator::Vitess')
-                        when 'none', nil then NoopPropagator.new
+                        when 'none', nil then nil
                         else
                           OpenTelemetry.logger.warn "The #{propagator} propagator is unknown and cannot be configured"
-                          NoopPropagator.new
                         end
         end
 
diff --git a/instrumentation/trilogy/lib/opentelemetry/instrumentation/trilogy/patches/client.rb b/instrumentation/trilogy/lib/opentelemetry/instrumentation/trilogy/patches/client.rb
index 1f7ae1fae..d5e93e7a8 100644
--- a/instrumentation/trilogy/lib/opentelemetry/instrumentation/trilogy/patches/client.rb
+++ b/instrumentation/trilogy/lib/opentelemetry/instrumentation/trilogy/patches/client.rb
@@ -50,7 +50,14 @@ def query(sql)
               ),
               kind: :client
             ) do |_span, context|
-              propagator.inject(sql, context: 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
diff --git a/instrumentation/trilogy/test/opentelemetry/instrumentation/trilogy/instrumentation_test.rb b/instrumentation/trilogy/test/opentelemetry/instrumentation/trilogy/instrumentation_test.rb
index 0bc3a6324..9b1b7b845 100644
--- a/instrumentation/trilogy/test/opentelemetry/instrumentation/trilogy/instrumentation_test.rb
+++ b/instrumentation/trilogy/test/opentelemetry/instrumentation/trilogy/instrumentation_test.rb
@@ -375,14 +375,52 @@
     describe 'when propagator is set to vitess' do
       let(:config) { { propagator: 'vitess' } }
 
-      it 'does inject context' do
+      it 'does inject context on frozen strings' do
+        sql = 'SELECT * from users where users.id = 1 and users.email = "test@test.com"'
+        assert(sql.frozen?)
+        propagator = OpenTelemetry::Instrumentation::Trilogy::Instrumentation.instance.propagator
+
+        arg_cache = {} # maintain handles to args
+        allow(client).to receive(:query).and_wrap_original do |m, *args|
+          arg_cache[:query_input] = args[0]
+          assert(args[0].frozen?)
+          m.call(args[0])
+        end
+
+        allow(propagator).to receive(:inject).and_wrap_original do |m, *args|
+          arg_cache[:inject_input] = args[0]
+          refute(args[0].frozen?)
+          assert_match(sql, args[0])
+          m.call(args[0], context: args[1][:context])
+        end
+
+        expect do
+          client.query(sql)
+        end.must_raise Trilogy::Error
+
+        # arg_cache[:inject_input] _was_ a mutable string, so it has the context injected
+        encoded = Base64.strict_encode64("{\"uber-trace-id\":\"#{span.hex_trace_id}:#{span.hex_span_id}:0:1\"}")
+        assert_equal(arg_cache[:inject_input], "/*VT_SPAN_CONTEXT=#{encoded}*/#{sql}")
+
+        # arg_cache[:inject_input] is now frozen
+        assert(arg_cache[:inject_input].frozen?)
+      end
+
+      it 'does inject context on unfrozen strings' do
+        # inbound SQL is not frozen (string prefixed with +)
         sql = +'SELECT * from users where users.id = 1 and users.email = "test@test.com"'
-        original_sql = sql.dup
+        refute(sql.frozen?)
+
+        # dup sql for comparison purposes, since propagator  mutates it
+        cached_sql = sql.dup
+
         expect do
           client.query(sql)
         end.must_raise Trilogy::Error
+
         encoded = Base64.strict_encode64("{\"uber-trace-id\":\"#{span.hex_trace_id}:#{span.hex_span_id}:0:1\"}")
-        _(sql).must_equal "/*VT_SPAN_CONTEXT=#{encoded}*/#{original_sql}"
+        assert_equal(sql, "/*VT_SPAN_CONTEXT=#{encoded}*/#{cached_sql}")
+        refute(sql.frozen?)
       end
     end