feat: collect pg db.collection_name attribute (#1087)

The db.collection.name attribute is conditionally required for Database spans. This PR uses regex to detect the collection name (called table name in PG) and report it as the db.collection.name attribute for PG instrumentation.

---------

Co-authored-by: Ariel Valentin <[email protected]>
Co-authored-by: Kayla Reopelle <[email protected]>
Co-authored-by: Robb Kidd <[email protected]>

instrumentation/pg/example/pg.rb

@@ -18,8 +18,17 @@
   password: ENV.fetch('TEST_POSTGRES_PASSWORD') { 'postgres' }
 )
 
-# Spans will be printed to your terminal when this statement executes:
+# Create a table
+conn.exec('CREATE TABLE test_table (id SERIAL PRIMARY KEY, name VARCHAR(50), age INT)')
+# Insert data into the table
+conn.exec("INSERT INTO test_table (name, age) VALUES ('Peter', 60), ('Paul', 25), ('Mary', 45)")
+
+# Spans will be printed to your terminal when these statements execute:
 conn.exec('SELECT 1 AS a, 2 AS b, NULL AS c').each_row { |r| puts r.inspect }
+conn.exec('SELECT * FROM test_table').each_row { |r| puts r.inspect }
+
+# Drop table when done querying
+conn.exec('DROP TABLE test_table')
 
 # You can use parameterized queries like so:
 # conn.exec_params('SELECT $1 AS a, $2 AS b, $3 AS c', [1, 2, nil]).each_row { |r| puts r.inspect }

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

@@ -14,6 +14,9 @@ module PG
       module Patches
         # Module to prepend to PG::Connection for instrumentation
         module Connection # rubocop:disable Metrics/ModuleLength
+          # 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
+
           PG::Constants::EXEC_ISH_METHODS.each do |method|
             define_method method do |*args, &block|
               span_name, attrs = span_attrs(:query, *args)
@@ -86,11 +89,13 @@ def lru_cache
           # module size limit! We can't win here unless we want to start
           # abstracting things into a million pieces.
           def span_attrs(kind, *args)
+            text = args[0]
+
             if kind == :query
-              operation = extract_operation(args[0])
-              sql = obfuscate_sql(args[0]).to_s
+              operation = extract_operation(text)
+              sql = obfuscate_sql(text).to_s
             else
-              statement_name = args[0]
+              statement_name = text
 
               if kind == :prepare
                 sql = obfuscate_sql(args[1]).to_s
@@ -104,6 +109,7 @@ def span_attrs(kind, *args)
 
             attrs = { 'db.operation' => validated_operation(operation), 'db.postgresql.prepared_statement_name' => statement_name }
             attrs['db.statement'] = sql unless config[:db_statement] == :omit
+            attrs['db.collection.name'] = collection_name(text)
             attrs.merge!(OpenTelemetry::Instrumentation::PG.attributes)
             attrs.compact!
 
@@ -125,6 +131,12 @@ def validated_operation(operation)
             operation if PG::Constants::SQL_COMMANDS.include?(operation)
           end
 
+          def collection_name(text)
+            text.scan(TABLE_NAME).flatten[0]
+          rescue StandardError
+            nil
+          end
+
           def client_attributes
             attributes = {
               'db.system' => 'postgresql',

instrumentation/pg/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/pg/test/opentelemetry/instrumentation/pg/instrumentation_test.rb

@@ -12,13 +12,16 @@
 require_relative '../../../../lib/opentelemetry/instrumentation/pg/patches/connection'
 
 # This test suite requires a running postgres container and dedicated test container
-# To run tests:
+# To run tests locally:
 # 1. Build the opentelemetry/opentelemetry-ruby-contrib image
 # - docker-compose build
 # 2. Bundle install
 # - docker-compose run ex-instrumentation-pg-test bundle install
-# 3. Run test suite
-# - docker-compose run ex-instrumentation-pg-test bundle exec rake test
+# 3. Install the dependencies for each Appraisal (https://github.com/thoughtbot/appraisal)
+# - docker-compose run ex-instrumentation-pg-test bundle exec appraisal install
+# 4. Run test suite with Appraisal
+# - docker-compose run ex-instrumentation-pg-test bundle exec appraisal rake test
+
 describe OpenTelemetry::Instrumentation::PG::Instrumentation do
   let(:instrumentation) { OpenTelemetry::Instrumentation::PG::Instrumentation.instance }
   let(:exporter) { EXPORTER }
@@ -33,6 +36,7 @@
   after do
     # Force re-install of instrumentation
     instrumentation.instance_variable_set(:@installed, false)
+    client&.close
   end
 
   describe 'tracing' do
@@ -76,7 +80,8 @@
           'db.operation' => 'PREPARE FOR SELECT 1',
           'db.postgresql.prepared_statement_name' => 'bar',
           'net.peer.ip' => '192.168.0.1',
-          'peer.service' => 'example:custom'
+          'peer.service' => 'example:custom',
+          'db.collection.name' => 'test_table'
         }
       end
 
@@ -263,6 +268,13 @@
       assert(!span.events.first.attributes['exception.stacktrace'].nil?)
     end
 
+    it 'extracts table name' do
+      client.query('CREATE TABLE test_table (personid int, name VARCHAR(50))')
+
+      _(span.attributes['db.collection.name']).must_equal 'test_table'
+      client.query('DROP TABLE test_table') # Drop table to avoid conflicts
+    end
+
     describe 'when db_statement is obfuscate' do
       let(:config) { { db_statement: :obfuscate } }
 
@@ -361,5 +373,23 @@
         _(span.attributes['net.peer.port']).must_equal port.to_i if PG.const_defined?(:DEF_PORT)
       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 unless ENV['OMIT_SERVICES']
 end