Mongo dbo demo

Gemfile

@@ -36,6 +36,7 @@ group :development, :test do
   gem 'simplecov-console', require: false, group: :test
   gem 'webmock' if RUBY_VERSION >= '2.0.0'
   gem 'base64' if RUBY_VERSION >= '3.4.0'
+  gem 'mongo', '>= 2.24.1'
 
   gem 'opentelemetry-propagator-b3'
   gem 'opentelemetry-test-helpers'

lib/solarwinds_apm/patch/tag_sql/sw_dbo_utils.rb

@@ -11,24 +11,33 @@ module SWODboUtils
         def self.annotate_span_and_sql(sql)
           return sql if sql.to_s.empty?
 
-          current_span = ::OpenTelemetry::Trace.current_span
+          ::OpenTelemetry::Trace.current_span
+          transparent = annotated_traceparent
 
           annotated_sql = ''
-          if current_span.context.trace_flags.sampled?
-            traceparent = SolarWindsAPM::Utils.traceparent_from_context(current_span.context)
-            annotated_traceparent = "/*traceparent='#{traceparent}'*/"
-            current_span.add_attributes({ 'sw.query_tag' => annotated_traceparent })
-            annotated_sql = "#{sql} #{annotated_traceparent}"
-          else
-            annotated_sql = sql
-          end
+          annotated_sql = if transparent.empty?
+                            sql
+                          else
+                            "#{sql} #{transparent}"
+                          end
 
           SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] annotated_sql: #{annotated_sql}" }
           annotated_sql
         rescue StandardError => e
           SolarWindsAPM.logger.error { "[#{self.class}/#{__method__}] Failed to annotated sql. Error: #{e.message}" }
           sql
         end
+
+        def self.annotated_traceparent
+          current_span = ::OpenTelemetry::Trace.current_span
+          if current_span.context.trace_flags.sampled?
+            traceparent = SolarWindsAPM::Utils.traceparent_from_context(current_span.context)
+            current_span.add_attributes({ 'sw.query_tag' => "/*traceparent='#{traceparent}'*/" })
+            "/*traceparent='#{traceparent}'*/"
+          else
+            ''
+          end
+        end
       end
     end
   end

lib/solarwinds_apm/patch/tag_sql/sw_mongo_patch.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+# © 2023 SolarWinds Worldwide, LLC. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at:http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.
+
+module SolarWindsAPM
+  module Patch
+    module TagSql
+      module SWOMongoPatch
+        # Annotate the final OP_MSG document rather than the operation spec.
+        #
+        # ConnectionBase#deliver is the shared boundary immediately before
+        # serialization, so this works independently of the driver's
+        # OpenTelemetry implementation and its internal tracer classes.
+        def deliver(message, context, options = {})
+          annotate_message_comment(message)
+          super
+        end
+        private :deliver
+
+        private
+
+        def annotate_message_comment(message)
+          return unless defined?(::Mongo::Protocol::Msg) &&
+                        message.is_a?(::Mongo::Protocol::Msg)
+
+          main_document = message.instance_variable_get(:@main_document)
+          return unless main_document
+
+          traceparent = ::SolarWindsAPM::Patch::TagSql::SWODboUtils.annotated_traceparent
+          return if traceparent.nil? || traceparent.empty?
+
+          key = comment_key(main_document)
+          main_document[key] = annotated_comment(main_document[key], traceparent)
+        end
+
+        def comment_key(main_document)
+          return 'comment' if main_document.key?('comment')
+          return :comment if main_document.key?(:comment)
+
+          'comment'
+        end
+
+        def annotated_comment(original, traceparent)
+          case original
+          when nil
+            traceparent
+          when String
+            original.empty? ? traceparent : "#{original}; #{traceparent}"
+          when Hash, BSON::Document
+            # Real document: add traceparent as a sibling key.
+            # Guard against clobbering a user key of the same name.
+            doc = BSON::Document.new(original)
+            if doc.key?('traceparent') || doc.key?(:traceparent)
+              doc['swo_traceparent'] = traceparent
+            else
+              doc['traceparent'] = traceparent
+            end
+            doc
+          else
+            # Scalar BSON value (ObjectId, Time, Integer, Float, true/false, ...).
+            # These have no fields to extend, so wrap them in a new document.
+            BSON::Document.new(
+              'swo_original_comment' => original,
+              'traceparent' => traceparent
+            )
+          end
+        end
+      end
+    end
+  end
+end
+
+# ConnectionBase#deliver exists before 2.22.0 and remains the common
+# serialization path in 2.23.0 and newer. Do not couple this patch to
+# either of the driver's OpenTelemetry integration implementations.
+if defined?(::Mongo::Server::ConnectionBase)
+  ::Mongo::Server::ConnectionBase.prepend(
+    ::SolarWindsAPM::Patch::TagSql::SWOMongoPatch
+  )
+end
+
+# Why use execute_with_span or not execute_with_span?
+#
+# execute_with_span only runs when the driver's OpenTelemetry tracing is enabled.
+# Look at the real entry point, Tracer#trace_operation in tracer.rb:72-76:
+# And enabled? is driven by OTEL_RUBY_INSTRUMENTATION_MONGODB_ENABLED (tracer.rb:42-46).
+# So if that env var isn't set to true/1/yes, OperationTracer#trace_operation is never called,
+# and therefore your prepended execute_with_span never fires — your traceparent silently disappears.

lib/solarwinds_apm/patch/tag_sql_patch.rb

@@ -9,3 +9,4 @@
 require_relative 'tag_sql/sw_dbo_utils'
 require_relative 'tag_sql/sw_mysql2_patch'
 require_relative 'tag_sql/sw_pg_patch'
+require_relative 'tag_sql/sw_mongo_patch'

test/patch/sw_mongo_patch_integrate_test.rb

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+# Copyright (c) 2023 SolarWinds, LLC.
+# All rights reserved.
+
+require 'minitest_helper'
+require 'mongo'
+require './lib/solarwinds_apm/config'
+require './lib/solarwinds_apm/opentelemetry'
+require './lib/solarwinds_apm/support/txn_name_manager'
+require './lib/solarwinds_apm/otel_config'
+require './lib/solarwinds_apm/api'
+require './lib/solarwinds_apm/support'
+require './lib/solarwinds_apm/patch/tag_sql/sw_dbo_utils'
+require './lib/solarwinds_apm/patch/tag_sql/sw_mongo_patch'
+
+# rubocop:disable Naming/MethodName
+module SolarWindsAPM
+  module Span
+    def self.createSpan(trans_name, domain, span_time, has_error); end
+  end
+end
+# rubocop:enable Naming/MethodName
+
+# ConnectionBase#deliver performs a live socket write, so exercise the real
+# patch module through a test double that returns the delivered message. The
+# mongo gem is installed for tests, so the real Mongo::Protocol::Msg is used.
+class FakeMongoConnectionBase
+  def deliver(message, _context, _options = {})
+    message
+  end
+  private :deliver
+end
+FakeMongoConnectionBase.prepend(SolarWindsAPM::Patch::TagSql::SWOMongoPatch)
+
+describe 'mongo patch integrate test' do
+  puts "\n\033[1m=== TEST RUN MONGO PATCH TEST: #{RUBY_VERSION} #{File.basename(__FILE__)} #{Time.now.strftime('%Y-%m-%d %H:%M')} ===\033[0m\n"
+
+  let(:sdk) { OpenTelemetry::SDK }
+  let(:exporter) { sdk::Trace::Export::InMemorySpanExporter.new }
+  let(:span_processor) { sdk::Trace::Export::SimpleSpanProcessor.new(exporter) }
+  let(:traceparent_pattern) { %r{/\*traceparent='[\da-f-]+'*\*/$} }
+
+  before do
+    OpenTelemetry::SDK.configure
+    OpenTelemetry.tracer_provider.add_span_processor(span_processor)
+    @tracer = OpenTelemetry.tracer_provider.tracer('mongo-patch-test')
+  end
+
+  def deliver_in_span(message)
+    @tracer.in_span('mongo-op') do |_span|
+      FakeMongoConnectionBase.new.send(:deliver, message, nil)
+    end
+    message.instance_variable_get(:@main_document)
+  end
+
+  it 'injects traceparent comment into the message main document and sets sw.query_tag attribute' do
+    message = Mongo::Protocol::Msg.new([], {}, {})
+
+    main_doc = deliver_in_span(message)
+
+    finished_spans = exporter.finished_spans
+    assert_match traceparent_pattern, finished_spans[0].attributes['sw.query_tag'], "Doesn't match sw.query_tag"
+    assert_match traceparent_pattern, main_doc['comment'], "message comment doesn't contain traceparent"
+  end
+
+  it 'appends traceparent comment when the message main document already has a string comment' do
+    message = Mongo::Protocol::Msg.new([], {}, { 'comment' => 'existing' })
+
+    main_doc = deliver_in_span(message)
+
+    assert_match %r{^existing; /\*traceparent='[\da-f-]+'*\*/$}, main_doc['comment'], "message comment doesn't append traceparent"
+  end
+
+  it 'appends traceparent comment when the main document uses a symbol comment key' do
+    message = Mongo::Protocol::Msg.new([], {}, { comment: 'existing' })
+
+    main_doc = deliver_in_span(message)
+
+    assert_match %r{^existing; /\*traceparent='[\da-f-]+'*\*/$}, main_doc[:comment], "message comment doesn't append traceparent"
+  end
+
+  it 'adds traceparent as a sibling key when the comment is a document' do
+    message = Mongo::Protocol::Msg.new([], {}, { 'comment' => { 'foo' => 'bar' } })
+
+    main_doc = deliver_in_span(message)
+
+    comment = main_doc['comment']
+    assert_instance_of BSON::Document, comment
+    assert_equal 'bar', comment['foo']
+    assert_match traceparent_pattern, comment['traceparent'], "document comment doesn't contain traceparent"
+  end
+
+  it 'preserves a user traceparent key by storing the trace under swo_traceparent' do
+    message = Mongo::Protocol::Msg.new([], {}, { 'comment' => { 'traceparent' => 'user-value' } })
+
+    main_doc = deliver_in_span(message)
+
+    comment = main_doc['comment']
+    assert_instance_of BSON::Document, comment
+    assert_equal 'user-value', comment['traceparent']
+    assert_match traceparent_pattern, comment['swo_traceparent'], "document comment doesn't preserve user traceparent"
+  end
+
+  it 'wraps a scalar comment in a document with the original value' do
+    message = Mongo::Protocol::Msg.new([], {}, { 'comment' => 42 })
+
+    main_doc = deliver_in_span(message)
+
+    comment = main_doc['comment']
+    assert_instance_of BSON::Document, comment
+    assert_equal 42, comment['swo_original_comment']
+    assert_match traceparent_pattern, comment['traceparent'], "scalar comment doesn't contain traceparent"
+  end
+end

test/patch/sw_mongo_patch_test.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+# Copyright (c) 2023 SolarWinds, LLC.
+# All rights reserved.
+
+require 'minitest_helper'
+require 'mongo'
+require './lib/solarwinds_apm/config'
+require './lib/solarwinds_apm/opentelemetry'
+require './lib/solarwinds_apm/support/txn_name_manager'
+require './lib/solarwinds_apm/otel_config'
+
+describe 'mongo patch test' do
+  puts "\n\033[1m=== TEST RUN MONGO PATCH TEST: #{RUBY_VERSION} #{File.basename(__FILE__)} #{Time.now.strftime('%Y-%m-%d %H:%M')} ===\033[0m\n"
+
+  it 'does not prepend the SWO mongo patch to Mongo::Server::ConnectionBase when tag_sql is false' do
+    SolarWindsAPM::Config[:tag_sql] = false
+    SolarWindsAPM::OTelConfig.initialize
+
+    connection_ancestors = Mongo::Server::ConnectionBase.ancestors.map(&:to_s)
+    refute_includes connection_ancestors, 'SolarWindsAPM::Patch::TagSql::SWOMongoPatch'
+  end
+end