feat: add config[:obfuscation_limit] to pg and mysql2

instrumentation/mysql2/lib/opentelemetry/instrumentation/mysql2/instrumentation.rb

@@ -22,6 +22,7 @@ class Instrumentation < OpenTelemetry::Instrumentation::Base
         option :peer_service, default: nil, validate: :string
         option :db_statement, default: :include, validate: %I[omit include obfuscate]
         option :span_name, default: :statement_type, validate: %I[statement_type db_name db_operation_and_name]
+        option :obfuscation_limit, default: 2000, validate: :integer
 
         private
 

instrumentation/mysql2/lib/opentelemetry/instrumentation/mysql2/patches/client.rb

@@ -30,6 +30,7 @@ module Client # rubocop:disable Metrics/ModuleLength
 
           QUERY_NAME_RE = Regexp.new("^(#{QUERY_NAMES.join('|')})", Regexp::IGNORECASE)
 
+          # From: https://github.com/newrelic/newrelic-ruby-agent/blob/0235b288d85b8bc795bdc1a24621dd9f84cfef45/lib/new_relic/agent/database/obfuscation_helpers.rb#L9-L34
           COMPONENTS_REGEX_MAP = {
             single_quotes: /'(?:[^']|'')*?(?:\\'.*|'(?!'))/,
             double_quotes: /"(?:[^"]|"")*?(?:\\".*|"(?!"))/,
@@ -70,8 +71,9 @@ def query(sql, options = {})
           private
 
           def obfuscate_sql(sql)
-            if sql.size > 2000
-              'SQL query too large to remove sensitive data ...'
+            if sql.size > config[:obfuscation_limit]
+              truncated_sql = sql[..sql.index(generated_mysql_regex) - 1]
+              truncated_sql + "...\nSQL truncated (> #{config[:obfuscation_limit]} characters)"
             else
               obfuscated = OpenTelemetry::Common::Utilities.utf8_encode(sql, binary: true)
               obfuscated = obfuscated.gsub(generated_mysql_regex, '?')

instrumentation/mysql2/test/opentelemetry/instrumentation/mysql2/instrumentation_test.rb

@@ -172,7 +172,6 @@
         _(span.attributes['net.peer.name']).must_equal host.to_s
         _(span.attributes['net.peer.port']).must_equal port.to_s
       end
-
       it 'encodes invalid byte sequences for db.statement' do
         # \255 is off-limits https://en.wikipedia.org/wiki/UTF-8#Codepage_layout
         sql = "SELECT * from users where users.id = 1 and users.email = '[email protected]\255'"
@@ -185,6 +184,20 @@
         _(span.name).must_equal 'mysql'
         _(span.attributes['db.statement']).must_equal obfuscated_sql
       end
+
+      describe 'with obfuscation_limit' do
+        let(:config) { { db_statement: :obfuscate, obfuscation_limit: 10 } }
+
+        it 'truncates SQL using config limit' do
+          sql = "SELECT * from users where users.id = 1 and users.email = '[email protected]'"
+          obfuscated_sql = "SELECT * from users where users.id = ...\nSQL truncated (> 10 characters)"
+          expect do
+            client.query(sql)
+          end.must_raise Mysql2::Error
+
+          _(span.attributes['db.statement']).must_equal obfuscated_sql
+        end
+      end
     end
 
     describe 'when db_statement set as omit' do

instrumentation/pg/lib/opentelemetry/instrumentation/pg/instrumentation.rb

@@ -26,6 +26,7 @@ class Instrumentation < OpenTelemetry::Instrumentation::Base
 
         option :peer_service, default: nil, validate: :string
         option :db_statement, default: :include, validate: %I[omit include obfuscate]
+        option :obfuscation_limit, default: 2000, validate: :integer
 
         private
 

instrumentation/pg/lib/opentelemetry/instrumentation/pg/patches/connection.rb

@@ -107,8 +107,10 @@ def validated_operation(operation)
           def obfuscate_sql(sql)
             return sql unless config[:db_statement] == :obfuscate
 
-            # Borrowed from opentelemetry-instrumentation-mysql2
-            return 'SQL query too large to remove sensitive data ...' if sql.size > 2000
+            if sql.size > config[:obfuscation_limit]
+              truncated_sql = sql[..sql.index(generated_postgres_regex) - 1]
+              return truncated_sql + "...\nSQL truncated (> #{config[:obfuscation_limit]} characters)"
+            end
 
             # From:
             # https://github.com/newrelic/newrelic-ruby-agent/blob/9787095d4b5b2d8fcaf2fdbd964ed07c731a8b6b/lib/new_relic/agent/database/obfuscator.rb

instrumentation/pg/test/opentelemetry/instrumentation/pg/instrumentation_test.rb

@@ -269,6 +269,20 @@
         _(span.attributes['net.peer.name']).must_equal host.to_s
         _(span.attributes['net.peer.port']).must_equal port.to_i
       end
+
+      describe 'with obfuscation_limit' do
+        let(:config) { { db_statement: :obfuscate, obfuscation_limit: 10 } }
+
+        it 'truncates SQL using config limit' do
+          sql = "SELECT * from users where users.id = 1 and users.email = '[email protected]'"
+          obfuscated_sql = "SELECT * from users where users.id = ...\nSQL truncated (> 10 characters)"
+          expect do
+            client.exec(sql)
+          end.must_raise PG::UndefinedTable
+
+          _(span.attributes['db.statement']).must_equal obfuscated_sql
+        end
+      end
     end
 
     describe 'when db_statement is omit' do