Skip to content

Commit

Permalink
fix: dup string if frozen in trilogy query (#863)
Browse files Browse the repository at this point in the history
* fix: propagate context from trilogy instrumentation for frozen strings
  • Loading branch information
plantfansam authored Feb 20, 2024
1 parent 6eaadb4 commit a2669ad
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 22 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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|
Expand Down Expand Up @@ -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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 = "[email protected]"'
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 = "[email protected]"'
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

Expand Down

0 comments on commit a2669ad

Please sign in to comment.