feat: Add with_attributes context propagation for PG instrumentation (#101)

* feat: Add `with_attributes` context propagation for PG instrumentation

* remove context_with_attributes

* add test cases

Co-authored-by: Andrew Hayworth <[email protected]>

Author: yoheyk

instrumentation/pg/README.md

@@ -30,6 +30,17 @@ OpenTelemetry::SDK.configure do |c|
 end
 ```
 
+The `PG` instrumentation allows the user to supply additional attributes via the `with_attributes` method. This makes it possible to supply additional attributes on PG spans. Attributes supplied in `with_attributes` supersede those automatically generated within `PG`'s automatic instrumentation. If you supply a `db.statement` attribute in `with_attributes`, this library's `:db_statement` configuration will not be applied.
+
+```ruby
+require 'opentelemetry/instrumentation/pg'
+
+conn = PG::Connection.open(host: "localhost", user: "root", dbname: "postgres")
+OpenTelemetry::Instrumentation::PG.with_attributes('pizzatoppings' => 'mushrooms') do
+  conn.exec("SELECT 1")
+end
+```
+
 ### Configuration options
 
 ```ruby

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

@@ -11,6 +11,35 @@ module OpenTelemetry
   module Instrumentation
     # Contains the OpenTelemetry instrumentation for the Pg gem
     module PG
+      extend self
+
+      CURRENT_ATTRIBUTES_KEY = Context.create_key('pg-attributes-hash')
+
+      private_constant :CURRENT_ATTRIBUTES_KEY
+
+      # Returns the attributes hash representing the postgres client context found
+      # in the optional context or the current context if none is provided.
+      #
+      # @param [optional Context] context The context to lookup the current
+      #   attributes hash. Defaults to Context.current
+      def attributes(context = nil)
+        context ||= Context.current
+        context.value(CURRENT_ATTRIBUTES_KEY) || {}
+      end
+
+      # Activates/deactivates the merged attributes hash within the current Context,
+      # which makes the "current attributes hash" available implicitly.
+      #
+      # On exit, the attributes hash that was active before calling this method
+      # will be reactivated.
+      #
+      # @param [Span] span the span to activate
+      # @yield [Hash, Context] yields attributes hash and a context containing the
+      #   attributes hash to the block.
+      def with_attributes(attributes_hash)
+        attributes_hash = attributes.merge(attributes_hash)
+        Context.with_value(CURRENT_ATTRIBUTES_KEY, attributes_hash) { |c, h| yield h, c }
+      end
     end
   end
 end

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

@@ -85,6 +85,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.merge!(OpenTelemetry::Instrumentation::PG.attributes)
             attrs.reject! { |_, v| v.nil? }
 
             [span_name(operation), client_attributes.merge(attrs)]

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

@@ -68,6 +68,42 @@
       _(span.attributes['peer.service']).must_equal 'readonly:postgres'
     end
 
+    describe '.attributes' do
+      let(:attributes) do
+        {
+          'db.name' => 'pg',
+          'db.statement' => 'foobar',
+          'db.operation' => 'PREPARE FOR SELECT 1',
+          'db.postgresql.prepared_statement_name' => 'bar',
+          'net.peer.ip' => '192.168.0.1',
+          'peer.service' => 'example:custom'
+        }
+      end
+
+      it 'returns an empty hash by default' do
+        _(OpenTelemetry::Instrumentation::PG.attributes).must_equal({})
+      end
+
+      it 'returns the current attributes hash' do
+        OpenTelemetry::Instrumentation::PG.with_attributes(attributes) do
+          _(OpenTelemetry::Instrumentation::PG.attributes).must_equal(attributes)
+        end
+      end
+
+      it 'sets span attributes according to with_attributes hash' do
+        OpenTelemetry::Instrumentation::PG.with_attributes(attributes) do
+          client.prepare('foo', 'SELECT 1')
+        end
+
+        _(span.attributes['db.name']).must_equal 'pg'
+        _(span.attributes['db.statement']).must_equal 'foobar'
+        _(span.attributes['db.operation']).must_equal 'PREPARE FOR SELECT 1'
+        _(span.attributes['db.postgresql.prepared_statement_name']).must_equal 'bar'
+        _(span.attributes['net.peer.ip']).must_equal '192.168.0.1'
+        _(span.attributes['peer.service']).must_equal 'example:custom'
+      end
+    end
+
     %i[exec query sync_exec async_exec].each do |method|
       it "after request (with method: #{method})" do
         client.send(method, 'SELECT 1')