Add Mysql2 and Trilogy collection_name attribute

Author: hannahramadan

instrumentation/mysql2/example/mysql2.rb

@@ -21,3 +21,6 @@
 client.query("SELECT * from information_schema.INNODB_TABLES; /**Dé**/").each do |row|
   puts row
 end
+
+client.query('CREATE TABLE test_table (id SERIAL PRIMARY KEY, name VARCHAR(50), age INT)')
+client.query('DROP TABLE test_table')

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

@@ -13,6 +13,10 @@ module Mysql2
       module Patches
         # Module to prepend to Mysql2::Client for instrumentation
         module Client
+
+          # Capture the first word (including letters, digits, underscores, & '.', ) that follows common table commands
+          TABLE_NAME = /\b(?:FROM|INTO|UPDATE|CREATE\s+TABLE(?:\s+IF\s+NOT\s+EXISTS)?|DROP\s+TABLE(?:\s+IF\s+EXISTS)?|ALTER\s+TABLE(?:\s+IF\s+EXISTS)?)\s+([\w\.]+)/i
+
           def query(sql, options = {})
             tracer.in_span(
               _otel_span_name(sql),
@@ -47,7 +51,7 @@ def _otel_span_name(sql)
           end
 
           def _otel_span_attributes(sql)
-            attributes = _otel_client_attributes
+            attributes = _otel_client_attributes(sql)
             case config[:db_statement]
             when :include
               attributes[SemanticConventions::Trace::DB_STATEMENT] = sql
@@ -68,7 +72,7 @@ def _otel_database_name
             (query_options[:database] || query_options[:dbname] || query_options[:db])&.to_s
           end
 
-          def _otel_client_attributes
+          def _otel_client_attributes(sql)
             # The client specific attributes can be found via the query_options instance variable
             # exposed on the mysql2 Client
             # https://github.com/brianmario/mysql2/blob/ca08712c6c8ea672df658bb25b931fea22555f27/lib/mysql2/client.rb#L25-L26
@@ -83,9 +87,18 @@ def _otel_client_attributes
 
             attributes[SemanticConventions::Trace::DB_NAME] = _otel_database_name
             attributes[SemanticConventions::Trace::PEER_SERVICE] = config[:peer_service]
+            attributes['db.collection.name'] = collection_name(sql)
+
             attributes
           end
 
+          def collection_name(sql)
+            sql.scan(TABLE_NAME).flatten[0]
+
+          rescue StandardError
+            nil
+          end
+
           def tracer
             Mysql2::Instrumentation.instance.tracer
           end

instrumentation/mysql2/test/fixtures/sql_table_name.json

@@ -0,0 +1,54 @@
+[
+    {
+      "name": "from",
+      "sql": "SELECT * FROM test_table"
+    },
+    {
+      "name": "select_count_from",
+      "sql": "SELECT COUNT(*) FROM test_table WHERE condition"
+    },
+    {
+      "name": "from_with_subquery",
+      "sql": "SELECT * FROM (SELECT * FROM test_table) AS table_alias"
+    },
+    {
+      "name": "insert_into",
+      "sql": "INSERT INTO test_table (column1, column2) VALUES (value1, value2)"
+    },
+    {
+      "name": "update",
+      "sql": "UPDATE test_table SET column1 = value1 WHERE condition"
+    },
+    {
+      "name": "delete_from",
+      "sql": "DELETE FROM test_table WHERE condition"
+    },
+    {
+      "name": "create_table",
+      "sql": "CREATE TABLE test_table (column1 datatype, column2 datatype)"
+    },
+    {
+      "name": "create_table_if_not_exists",
+      "sql": "CREATE TABLE IF NOT EXISTS test_table (column1 datatype, column2 datatype)"
+    },
+    {
+      "name": "alter_table",
+      "sql": "ALTER TABLE test_table ADD column_name datatype"
+    },
+    {
+      "name": "drop_table",
+      "sql": "DROP TABLE test_table"
+    },
+    {
+      "name": "drop_table_if_exists",
+      "sql": "DROP TABLE IF EXISTS test_table"
+    },
+    {
+      "name": "insert_into",
+      "sql": "INSERT INTO test_table values('', 'a''b c',0, 1 , 'd''e f''s h')"
+    },
+    {
+      "name": "from_with_join",
+      "sql": "SELECT columns FROM test_table JOIN table2 ON test_table.column = table2.column"
+    }
+  ]

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

@@ -15,9 +15,9 @@
 # 1. Build the opentelemetry/opentelemetry-ruby-contrib image
 # - docker-compose build
 # 2. Bundle install
-# - docker-compose run ex-instrumentation-mysql2-test bundle install
+# - docker-compose run ex-instrumentation-mysql2-test bundle exec appraisal install
 # 3. Run test suite
-# - docker-compose run ex-instrumentation-mysql2-test bundle exec rake test
+# - docker-compose run ex-instrumentation-mysql2-test bundle exec appraisal rake test
 describe OpenTelemetry::Instrumentation::Mysql2::Instrumentation do
   let(:instrumentation) { OpenTelemetry::Instrumentation::Mysql2::Instrumentation.instance }
   let(:exporter) { EXPORTER }
@@ -473,6 +473,25 @@
           end
         end
       end
+
+      describe '#connection_name' do
+
+        def self.load_fixture
+          data = File.read("#{Dir.pwd}/test/fixtures/sql_table_name.json")
+          JSON.parse(data)
+        end
+
+        load_fixture.each do |test_case|
+          name = test_case['name']
+          query = test_case['sql']
+
+          it "returns the table name for #{name}" do
+            table_name = client.send(:collection_name, query)
+
+            expect(table_name).must_equal('test_table')
+          end
+        end
+      end
     end
   end unless ENV['OMIT_SERVICES']
 end

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

@@ -13,6 +13,10 @@ module Trilogy
       module Patches
         # Module to prepend to Trilogy for instrumentation
         module Client
+
+          # Capture the first word (including letters, digits, underscores, & '.', ) that follows common table commands
+          TABLE_NAME = /\b(?:FROM|INTO|UPDATE|CREATE\s+TABLE(?:\s+IF\s+NOT\s+EXISTS)?|DROP\s+TABLE(?:\s+IF\s+EXISTS)?|ALTER\s+TABLE(?:\s+IF\s+EXISTS)?)\s+([\w\.]+)/i
+
           def initialize(options = {})
             @connection_options = options # This is normally done by Trilogy#initialize
 
@@ -76,6 +80,8 @@ def client_attributes(sql = nil)
             attributes['db.instance.id'] = @connected_host unless @connected_host.nil?
 
             if sql
+              attributes['db.collection.name'] = collection_name(sql)
+
               case config[:db_statement]
               when :obfuscate
                 attributes[::OpenTelemetry::SemanticConventions::Trace::DB_STATEMENT] =
@@ -84,10 +90,17 @@ def client_attributes(sql = nil)
                 attributes[::OpenTelemetry::SemanticConventions::Trace::DB_STATEMENT] = sql
               end
             end
-
+            attributes.compact!
             attributes
           end
 
+          def collection_name(sql)
+            sql.scan(TABLE_NAME).flatten[0]
+
+          rescue StandardError
+            nil
+          end
+
           def database_name
             connection_options[:database]
           end

instrumentation/trilogy/test/fixtures/sql_table_name.json

@@ -0,0 +1,54 @@
+[
+    {
+      "name": "from",
+      "sql": "SELECT * FROM test_table"
+    },
+    {
+      "name": "select_count_from",
+      "sql": "SELECT COUNT(*) FROM test_table WHERE condition"
+    },
+    {
+      "name": "from_with_subquery",
+      "sql": "SELECT * FROM (SELECT * FROM test_table) AS table_alias"
+    },
+    {
+      "name": "insert_into",
+      "sql": "INSERT INTO test_table (column1, column2) VALUES (value1, value2)"
+    },
+    {
+      "name": "update",
+      "sql": "UPDATE test_table SET column1 = value1 WHERE condition"
+    },
+    {
+      "name": "delete_from",
+      "sql": "DELETE FROM test_table WHERE condition"
+    },
+    {
+      "name": "create_table",
+      "sql": "CREATE TABLE test_table (column1 datatype, column2 datatype)"
+    },
+    {
+      "name": "create_table_if_not_exists",
+      "sql": "CREATE TABLE IF NOT EXISTS test_table (column1 datatype, column2 datatype)"
+    },
+    {
+      "name": "alter_table",
+      "sql": "ALTER TABLE test_table ADD column_name datatype"
+    },
+    {
+      "name": "drop_table",
+      "sql": "DROP TABLE test_table"
+    },
+    {
+      "name": "drop_table_if_exists",
+      "sql": "DROP TABLE IF EXISTS test_table"
+    },
+    {
+      "name": "insert_into",
+      "sql": "INSERT INTO test_table values('', 'a''b c',0, 1 , 'd''e f''s h')"
+    },
+    {
+      "name": "from_with_join",
+      "sql": "SELECT columns FROM test_table JOIN table2 ON test_table.column = table2.column"
+    }
+  ]

instrumentation/trilogy/test/opentelemetry/instrumentation/trilogy/instrumentation_test.rb

@@ -9,6 +9,16 @@
 require_relative '../../../../lib/opentelemetry/instrumentation/trilogy'
 require_relative '../../../../lib/opentelemetry/instrumentation/trilogy/patches/client'
 
+# This test suite requires a running mysql container and dedicated test container. We can use the same
+# docker-compose file as the mysql2 instrumentation tests. The test container should have the mysql client.
+# To run tests locally:
+# 1. Build the opentelemetry/opentelemetry-ruby-contrib image
+# - docker-compose build
+# 2. Open a bash shell in the test container and cd to the trilogy directory
+# - docker-compose run ex-instrumentation-mysql2-test bash -c 'cd ../trilogy && bash'
+# 3. Bundle install and run tests with the Appraisals gem
+# - bundle exec appraisal install && bundle exec appraisal rake test
+
 describe OpenTelemetry::Instrumentation::Trilogy do
   let(:instrumentation) { OpenTelemetry::Instrumentation::Trilogy::Instrumentation.instance }
   let(:exporter) { EXPORTER }
@@ -626,5 +636,23 @@
         end
       end
     end
+
+    describe '#connection_name' do
+      def self.load_fixture
+        data = File.read("#{Dir.pwd}/test/fixtures/sql_table_name.json")
+        JSON.parse(data)
+      end
+
+      load_fixture.each do |test_case|
+        name = test_case['name']
+        query = test_case['sql']
+
+        it "returns the table name for #{name}" do
+          table_name = client.send(:collection_name, query)
+
+          expect(table_name).must_equal('test_table')
+        end
+      end
+    end
   end
 end