feat: collect db.collection_name attribute - pg instrumentation

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 statement 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

@@ -86,11 +86,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 +106,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 +128,13 @@ def validated_operation(operation)
             operation if PG::Constants::SQL_COMMANDS.include?(operation)
           end
 
+          def collection_name(text)
+            # Capture the first word (including letters, digits, underscores, & '.', ) that follows common table commands
+            pattern = /\b(?:FROM|INTO|UPDATE|CREATE\s+TABLE(?:\s+IF\s+NOT\s+EXISTS)?|DROP\s+TABLE\s+IF\s+EXISTS)\s+([\w\.]+)/i
+
+            text.scan(pattern).flatten[0]
+          end
+
           def client_attributes
             attributes = {
               'db.system' => 'postgresql',

instrumentation/pg/test/fixtures/sql_table_name.json

@@ -0,0 +1,58 @@
+[
+    {
+      "name": "from",
+      "sql": "SELECT * FROM test_table"
+    },
+    {
+      "name": "select_count_from",
+      "sql": "SELECT COUNT(*) FROM table_name WHERE condition"
+    },
+    {
+      "name": "from_with_subquery",
+      "sql": "SELECT * FROM (SELECT * FROM table_name) AS table_alias"
+    },
+    {
+      "name": "insert_into",
+      "sql": "INSERT INTO table_name (column1, column2) VALUES (value1, value2)"
+    },
+    {
+      "name": "drop_table",
+      "sql": "DROP TABLE table_name"
+    },
+    {
+      "name": "update",
+      "sql": "UPDATE table_name SET column1 = value1 WHERE condition"
+    },
+    {
+      "name": "delete_from",
+      "sql": "DELETE FROM table_name WHERE condition"
+    },
+    {
+      "name": "create_table",
+      "sql": "CREATE TABLE table_name (column1 datatype, column2 datatype)"
+    },
+    {
+      "name": "create_table_if_not_exists",
+      "sql": "CREATE TABLE IF NOT EXISTS table_name (column1 datatype, column2 datatype)"
+    },
+    {
+      "name": "alter_table",
+      "sql": "ALTER TABLE table_name ADD column_name datatype"
+    },
+    {
+      "name": "drop_table",
+      "sql": "DROP TABLE table_name"
+    },
+    {
+      "name": "drop_table_if_exists",
+      "sql": "DROP TABLE IF EXISTS table_name"
+    },
+    {
+      "name": "insert_into",
+      "sql": "INSERT INTO X values('', 'a''b c',0, 1 , 'd''e f''s h')"
+    },
+    {
+      "name": "from_with_join",
+      "sql": "SELECT columns FROM table1 JOIN table2 ON table1.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,21 @@
         _(span.attributes['net.peer.port']).must_equal port.to_i if PG.const_defined?(:DEF_PORT)
       end
     end
+
+    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']
+
+      define_method(:"test_sql_table_name_#{name}") do
+        table_name = client.send(:collection_name, query)
+
+        assert('test_table', table_name)
+      end
+    end
   end unless ENV['OMIT_SERVICES']
 end