Skip to content

Commit

Permalink
feat: collect pg db.collection_name attribute (#1087)
Browse files Browse the repository at this point in the history
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 <arielvalentin@users.noreply.github.com>
Co-authored-by: Kayla Reopelle <87386821+kaylareopelle@users.noreply.github.com>
Co-authored-by: Robb Kidd <robbkidd@honeycomb.io>
  • Loading branch information
4 people authored Aug 14, 2024
1 parent cac6b47 commit 70fdad6
Show file tree
Hide file tree
Showing 4 changed files with 113 additions and 8 deletions.
11 changes: 10 additions & 1 deletion instrumentation/pg/example/pg.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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!

Expand All @@ -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',
Expand Down
54 changes: 54 additions & 0 deletions instrumentation/pg/test/fixtures/sql_table_name.json
Original file line number Diff line number Diff line change
@@ -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"
}
]
Original file line number Diff line number Diff line change
Expand Up @@ -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 }
Expand All @@ -33,6 +36,7 @@
after do
# Force re-install of instrumentation
instrumentation.instance_variable_set(:@installed, false)
client&.close
end

describe 'tracing' do
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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 } }

Expand Down Expand Up @@ -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

0 comments on commit 70fdad6

Please sign in to comment.