feat: Instrument PG connect (#1763)

* feat: add instrumentation for connecting to PG

Author: kirs

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

@@ -40,6 +40,7 @@ def require_dependencies
 
         def patch_client
           ::PG::Connection.prepend(Patches::Connection)
+          ::PG::Connection.singleton_class.prepend(Patches::ClassMethods)
         end
       end
     end

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

@@ -12,6 +12,53 @@ module OpenTelemetry
   module Instrumentation
     module PG
       module Patches
+        # Module to prepend to PG::Connection singleton class for connection initialization
+        module ClassMethods
+          def new(*args, &block)
+            tracer = PG::Instrumentation.instance.tracer
+            config = PG::Instrumentation.instance.config
+
+            tracer.in_span('connect', kind: :client) do |span|
+              conn = super(*args) do |yielded_conn|g
+                set_connection_attributes(span, yielded_conn, config)
+                block.call(yielded_conn) if block
+              end
+
+              unless block
+                set_connection_attributes(span, conn, config)
+              end
+
+              conn
+            end
+          end
+
+          private
+
+          def set_connection_attributes(span, conn, config)
+            attributes = {
+              'db.system' => 'postgresql',
+              'db.name' => conn.db,
+              'db.user' => conn.user
+            }
+            attributes['peer.service'] = config[:peer_service] if config[:peer_service]
+
+            h = conn.host
+            if h&.start_with?('/')
+              attributes['net.sock.family'] = 'unix'
+              attributes['net.peer.name'] = h
+            else
+              attributes['net.transport'] = 'ip_tcp'
+              attributes['net.peer.name'] = h
+              attributes['net.peer.port'] = conn.port if defined?(::PG::DEF_PGPORT)
+            end
+
+            attributes.merge!(OpenTelemetry::Instrumentation::PG.attributes)
+            attributes.compact!
+
+            span.add_attributes(attributes)
+          end
+        end
+
         # 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

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

@@ -64,6 +64,114 @@
       _(exporter.finished_spans.size).must_equal 0
     end
 
+    describe 'connection initialization' do
+      it 'creates a connect span when establishing a connection' do
+        conn = PG::Connection.open(
+          host: host,
+          port: port,
+          user: user,
+          dbname: dbname,
+          password: password
+        )
+        conn.close
+
+        _(exporter.finished_spans.size).must_equal 1
+        connect_span = exporter.finished_spans.first
+        _(connect_span.name).must_equal 'connect'
+        _(connect_span.kind).must_equal :client
+        _(connect_span.attributes['db.system']).must_equal 'postgresql'
+        _(connect_span.attributes['db.name']).must_equal dbname
+        _(connect_span.attributes['db.user']).must_equal user
+        _(connect_span.attributes['net.peer.name']).must_equal host
+        _(connect_span.attributes['net.transport']).must_equal 'ip_tcp'
+      end
+
+      it 'creates a connect span using PG::Connection.new' do
+        conn = PG::Connection.new(
+          host: host,
+          port: port,
+          user: user,
+          dbname: dbname,
+          password: password
+        )
+        conn.close
+
+        connect_span = exporter.finished_spans.first
+        _(connect_span.name).must_equal 'connect'
+        _(connect_span.kind).must_equal :client
+        _(connect_span.attributes['db.system']).must_equal 'postgresql'
+        _(connect_span.attributes['db.name']).must_equal dbname
+      end
+
+      it 'creates a connect span using PG.connect' do
+        conn = PG.connect(
+          host: host,
+          port: port,
+          user: user,
+          dbname: dbname,
+          password: password
+        )
+        conn.close
+
+        connect_span = exporter.finished_spans.first
+        _(connect_span.name).must_equal 'connect'
+        _(connect_span.kind).must_equal :client
+        _(connect_span.attributes['db.system']).must_equal 'postgresql'
+      end
+
+      it 'accepts peer service name from config for connection' do
+        instrumentation.instance_variable_set(:@installed, false)
+        instrumentation.install(peer_service: 'readonly:postgres')
+
+        conn = PG::Connection.open(
+          host: host,
+          port: port,
+          user: user,
+          dbname: dbname,
+          password: password
+        )
+        conn.close
+
+        connect_span = exporter.finished_spans.first
+        _(connect_span.attributes['peer.service']).must_equal 'readonly:postgres'
+      end
+
+      it 'sets unix socket attributes when connecting via socket' do
+        skip 'socket test requires no host/port' if host || port
+
+        conn = PG::Connection.open(
+          user: user,
+          dbname: dbname,
+          password: password
+        )
+        conn.close
+
+        connect_span = exporter.finished_spans.first
+        _(connect_span.name).must_equal 'connect'
+        _(connect_span.attributes['db.system']).must_equal 'postgresql'
+        _(connect_span.attributes['net.sock.family']).must_equal 'unix'
+        _(connect_span.attributes['net.peer.name']).must_match %r{^/}
+      end
+
+      it 'records connection errors' do
+        expect do
+          PG::Connection.open(
+            host: host,
+            port: port,
+            user: 'invalid_user',
+            dbname: dbname,
+            password: 'wrong_password'
+          )
+        end.must_raise PG::ConnectionBad
+
+        connect_span = exporter.finished_spans.first
+        _(connect_span.name).must_equal 'connect'
+        _(connect_span.status.code).must_equal OpenTelemetry::Trace::Status::ERROR
+        _(connect_span.events.first.name).must_equal 'exception'
+        _(connect_span.events.first.attributes['exception.type']).must_equal 'PG::ConnectionBad'
+      end
+    end
+
     it 'accepts peer service name from config' do
       instrumentation.instance_variable_set(:@installed, false)
       instrumentation.install(peer_service: 'readonly:postgres')