NH-89218: patch on db adapter for obfuscation and tracecontext in query

lib/solarwinds_apm.rb

@@ -20,7 +20,7 @@
 
   begin
     if RUBY_PLATFORM.include?('linux')
-      require 'solarwinds_apm/config'
+      require 'solarwinds_apm/config'            # initialize is called at end of config.rb file
       require 'solarwinds_apm/oboe_init_options' # setup oboe reporter options
       if !SolarWindsAPM::OboeInitOptions.instance.service_key_ok? && !SolarWindsAPM::OboeInitOptions.instance.lambda_env
         SolarWindsAPM.logger.warn '=============================================================='
@@ -60,7 +60,6 @@
           require 'solarwinds_apm/api'
           require 'solarwinds_apm/support'
           require 'solarwinds_apm/opentelemetry'
-          require 'solarwinds_apm/patch'
           require 'solarwinds_apm/otel_config'
 
           if ENV['SW_APM_AUTO_CONFIGURE'] != 'false'

lib/solarwinds_apm/opentelemetry/otlp_processor.rb

@@ -22,7 +22,7 @@ def initialize
       # @param [Context] parent_context the
       #  started span.
       def on_start(span, parent_context)
-        SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] processor on_start span: #{span.inspect}" }
+        SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] processor on_start span: #{span.to_span_data.inspect}" }
 
         return if non_entry_span(parent_context: parent_context)
 

lib/solarwinds_apm/opentelemetry/solarwinds_processor.rb

@@ -30,7 +30,7 @@ def initialize(txn_manager)
       #  started span.
       def on_start(span, parent_context)
         SolarWindsAPM.logger.debug do
-          "[#{self.class}/#{__method__}] processor on_start span: #{span.inspect}, parent_context: #{parent_context.inspect}"
+          "[#{self.class}/#{__method__}] processor on_start span: #{span.to_span_data.inspect}"
         end
 
         return if non_entry_span(parent_context: parent_context)
@@ -47,7 +47,7 @@ def on_start(span, parent_context)
       #
       # @param [Span] span the {Span} that just ended.
       def on_finish(span)
-        SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] processor on_finish span: #{span.inspect}" }
+        SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] processor on_finish span: #{span.to_span_data.inspect}" }
 
         return if non_entry_span(span: span)
 

lib/solarwinds_apm/opentelemetry/solarwinds_propagator.rb

@@ -54,7 +54,7 @@ def inject(carrier, context: ::OpenTelemetry::Context.current,
           SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] inject context: #{context.inspect}" }
 
           span_context = ::OpenTelemetry::Trace.current_span(context)&.context
-          SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] span_context #{span_context.inspect}" }
+          SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] span_context: #{span_context.inspect}" }
           return unless span_context&.valid?
 
           trace_flag = span_context.trace_flags.sampled? ? 1 : 0

lib/solarwinds_apm/opentelemetry/solarwinds_response_propagator.rb

@@ -35,9 +35,10 @@ def extract(carrier, context: ::OpenTelemetry::Context.current, getter: ::OpenTe
         def inject(carrier, context: ::OpenTelemetry::Context.current,
                    setter: ::OpenTelemetry::Context::Propagation.text_map_setter)
           span_context = ::OpenTelemetry::Trace.current_span(context).context
-          SolarWindsAPM.logger.debug do
-            "[#{self.class}/#{__method__}] context: #{context.inspect}; span_context: #{span_context.inspect}"
-          end
+
+          SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] context current_span: #{context.instance_variable_get(:@entries)&.values&.first.inspect}" }
+          SolarWindsAPM.logger.debug { "[#{self.class}/#{__method__}] span_context: #{span_context.inspect}" }
+
           return unless span_context&.valid?
 
           x_trace                = Utils.traceparent_from_context(span_context)

lib/solarwinds_apm/otel_config.rb

@@ -126,6 +126,15 @@ def self.initialize
 
       # configure sampler afterwards
       ::OpenTelemetry.tracer_provider.sampler = @@config[:sampler]
+
+      require_relative 'patches/tag_sql_patch' if SolarWindsAPM::Config[:tag_sql]
+
+      if ENV['SW_APM_AUTO_CONFIGURE'] == 'false'
+        SolarWindsAPM.logger.info '==================================================================='
+        SolarWindsAPM.logger.info "\e[1mSolarWindsAPM manual initialization was successful.\e[0m"
+        SolarWindsAPM.logger.info '==================================================================='
+      end
+
       nil
     end
 

lib/solarwinds_apm/patches/mysql2_client_patch.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+module SolarWindsAPM
+  module Patches
+    module SWOMysql2ClientPatch
+      # this method is just save the traceparent comment before sanitization
+      def _otel_span_attributes(sql)
+        # if omit, then no need to append any statement
+        # if include, then comments won't be removed
+        # only obfuscate need to add the original comments back
+        #   because of obfuscation method removed the comments
+        # This module need to injected after Mysql2::Patches::Client prepended (aka after solarwinds_apm initialized)
+        # check ::Mysql2::Client.included_modules or ::Mysql2::Client.ancestors to see the order of calling
+        if config[:db_statement] == :obfuscate
+          extracted_comments = sql.match(TagSqlConstants::TRACEPARENT_REGEX)
+          attributes = super
+          attributes[::OpenTelemetry::SemanticConventions::Trace::DB_STATEMENT] = attributes[::OpenTelemetry::SemanticConventions::Trace::DB_STATEMENT] + extracted_comments&.match(0).to_s
+          attributes
+        else
+          super
+        end
+      end
+    end
+  end
+end
+
+Mysql2::Client.prepend(SolarWindsAPM::Patches::SWOMysql2ClientPatch) if defined?(Mysql2::Client) && defined?(OpenTelemetry::Instrumentation::Mysql2::Patches::Client)

lib/solarwinds_apm/patches/pg_connection_patch.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+module SolarWindsAPM
+  module Patches
+    module SWOPgConnectionPatch
+      # this method is just save the traceparent comment before sanitization
+      def obfuscate_sql(sql)
+        if config[:db_statement] == :obfuscate
+          extracted_comments = sql.match(TagSqlConstants::TRACEPARENT_REGEX)
+          super + extracted_comments&.match(0).to_s
+        else
+          super
+        end
+      end
+    end
+  end
+end
+
+PG::Connection.prepend(SolarWindsAPM::Patches::SWOPgConnectionPatch) if defined?(PG::Connection) && defined?(OpenTelemetry::Instrumentation::PG::Patches::Connection)

lib/solarwinds_apm/patches/tag_sql_constants.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+module SolarWindsAPM
+  module Patches
+    module TagSqlConstants
+      TRACEPARENT_REGEX = %r{/\*\s*traceparent=?'?[0-9a-f]{2}-[0-9a-f]{32}-[0-9a-f]{16}-[0-9a-f]{2}'?\s*\*/}
+    end
+  end
+end

lib/solarwinds_apm/patches/tag_sql_patch.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require_relative 'tag_sql_constants'
+require_relative 'mysql2_client_patch'
+require_relative 'pg_connection_patch'
+require_relative 'trilogy_client_patch'

lib/solarwinds_apm/patches/trilogy_client_patch.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+module SolarWindsAPM
+  module Patches
+    module SWOTrilogyClientPatch
+      # this method is just save the traceparent comment before sanitization
+      def client_attributes(sql = nil)
+        if sql && config[:db_statement] == :obfuscate
+          extracted_comments = sql.match(TagSqlConstants::TRACEPARENT_REGEX)
+          attributes = super
+          attributes[::OpenTelemetry::SemanticConventions::Trace::DB_STATEMENT] = attributes[::OpenTelemetry::SemanticConventions::Trace::DB_STATEMENT] + extracted_comments&.match(0).to_s
+          attributes
+        else
+          super
+        end
+      end
+    end
+  end
+end
+
+Trilogy.prepend(SolarWindsAPM::Patches::SWOTrilogyClientPatch) if defined?(Trilogy) && defined?(OpenTelemetry::Instrumentation::Trilogy::Patches::Client)

lib/solarwinds_apm/support/swomarginalia/README.md

@@ -33,6 +33,10 @@ This folder contains the code that is copied from original [marginalia](https://
 
 1. (new file) prepend the ActiveRecordInstrumentation to activerecord adapter
 
+### annotation.rb
+1. (new file) utilize comment file append the generated traceparent data into sql
+2. If sql already contains the same comments, then it won't inject again
+
 ## Example
 
 ### Sample output of rails application

lib/solarwinds_apm/support/swomarginalia/annotation.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require_relative 'comment'
+
+module SolarWindsAPM
+  module SWOMarginalia
+    module Annotation
+      def annotate_sql(sql)
+        SWOMarginalia::Comment.update_adapter!(self)            # switch to current sql adapter
+        comment = SWOMarginalia::Comment.construct_comment      # comment will include traceparent
+        if comment.present? && !sql.include?(comment)
+          sql = if SWOMarginalia::Comment.prepend_comment
+                  "/*#{comment}*/ #{sql}"
+                else
+                  "#{sql} /*#{comment}*/"
+                end
+        end
+
+        inline_comment = SWOMarginalia::Comment.construct_inline_comment # this is for customized_swo_inline_annotations (user-defined value)
+        if inline_comment.present? && !sql.include?(inline_comment)
+          sql = if SWOMarginalia::Comment.prepend_comment
+                  "/*#{inline_comment}*/ #{sql}"
+                else
+                  "#{sql} /*#{inline_comment}*/"
+                end
+        end
+
+        sql
+      end
+
+      # We don't want to trace framework caches.
+      # Only instrument SQL that directly hits the database.
+      def ignore_payload?(name)
+        %w[SCHEMA EXPLAIN CACHE].include?(name.to_s)
+      end
+    end
+  end
+end

lib/solarwinds_apm/support/swomarginalia/load_swomarginalia.rb

@@ -41,14 +41,10 @@ def self.insert_into_action_controller
       end
 
       def self.insert_into_active_record
-        ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(SWOMarginalia::ActiveRecordInstrumentation) if defined? ActiveRecord::ConnectionAdapters::Mysql2Adapter
-        ActiveRecord::ConnectionAdapters::MysqlAdapter.prepend(SWOMarginalia::ActiveRecordInstrumentation) if defined? ActiveRecord::ConnectionAdapters::MysqlAdapter
+        ActiveRecord::ConnectionAdapters::Mysql2Adapter.prepend(SWOMarginalia::ActiveRecordInstrumentation)     if defined? ActiveRecord::ConnectionAdapters::Mysql2Adapter
         ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(SWOMarginalia::ActiveRecordInstrumentation) if defined? ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
-        return unless defined? ActiveRecord::ConnectionAdapters::SQLite3Adapter
-
-        return unless defined? ActiveRecord::ConnectionAdapters::SQLite3Adapter
-
-        ActiveRecord::ConnectionAdapters::SQLite3Adapter.prepend(SWOMarginalia::ActiveRecordInstrumentation)
+        ActiveRecord::ConnectionAdapters::SQLite3Adapter.prepend(SWOMarginalia::ActiveRecordInstrumentation)    if defined? ActiveRecord::ConnectionAdapters::SQLite3Adapter
+        ActiveRecord::ConnectionAdapters::TrilogyAdapter.prepend(SWOMarginalia::ActiveRecordInstrumentation)    if defined? ActiveRecord::ConnectionAdapters::TrilogyAdapter
       end
     end
   end

lib/solarwinds_apm/support/swomarginalia/swomarginalia.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require_relative 'comment'
+require_relative 'annotation'
 
 module SolarWindsAPM
   module SWOMarginalia
@@ -9,6 +9,8 @@ module SWOMarginalia
     # This ActiveRecordInstrumentation should only work for activerecord < 7.0 since after rails 7
     # this module won't be prepend to activerecord
     module ActiveRecordInstrumentation
+      include Annotation
+
       def execute(sql, *args, **options)
         super(annotate_sql(sql), *args, **options)
       end
@@ -18,46 +20,25 @@ def execute_and_clear(sql, *args, **options)
         super(annotate_sql(sql), *args, **options)
       end
 
+      # For activerecord < 7.1.0
       def exec_query(sql, *args, **options)
         super(annotate_sql(sql), *args, **options)
       end
 
+      # This patch is for non-rails app (e.g. sinatra) that use activerecord >= 7.1.0
+      # From 7.1.0, query() changed calling from exec_query() to internal_exec_query()
+      # Psql is not affected because in psql, internal_exec_query calls execute_and_clear
+      def internal_exec_query(sql, *args, **options)
+        super(annotate_sql(sql), *args, **options)
+      end
+
       def exec_delete(sql, *args)
         super(annotate_sql(sql), *args)
       end
 
       def exec_update(sql, *args)
         super(annotate_sql(sql), *args)
       end
-
-      def annotate_sql(sql)
-        SWOMarginalia::Comment.update_adapter!(self)            # switch to current sql adapter
-        comment = SWOMarginalia::Comment.construct_comment      # comment will include traceparent
-        if comment.present? && !sql.include?(comment)
-          sql = if SWOMarginalia::Comment.prepend_comment
-                  "/*#{comment}*/ #{sql}"
-                else
-                  "#{sql} /*#{comment}*/"
-                end
-        end
-
-        inline_comment = SWOMarginalia::Comment.construct_inline_comment # this is for customized_swo_inline_annotations (user-defined value)
-        if inline_comment.present? && !sql.include?(inline_comment)
-          sql = if SWOMarginalia::Comment.prepend_comment
-                  "/*#{inline_comment}*/ #{sql}"
-                else
-                  "#{sql} /*#{inline_comment}*/"
-                end
-        end
-
-        sql
-      end
-
-      # We don't want to trace framework caches.
-      # Only instrument SQL that directly hits the database.
-      def ignore_payload?(name)
-        %w[SCHEMA EXPLAIN CACHE].include?(name.to_s)
-      end
     end
 
     module ActionControllerInstrumentation

test/minitest_helper.rb

@@ -359,3 +359,15 @@ def clean_old_setting
   ENV.delete('OTEL_TRACES_EXPORTER')
   ENV.delete('SW_APM_ENABLED')
 end
+
+##
+# Some Marco
+#
+
+module OpenTelemetry
+  module SemanticConventions
+    module Trace
+      DB_STATEMENT = 'db.statement'
+    end
+  end
+end