lib/opentelemetry/instrumentation/pg/patches/connection.rb in opentelemetry-instrumentation-pg-0.25.2 vs lib/opentelemetry/instrumentation/pg/patches/connection.rb in opentelemetry-instrumentation-pg-0.25.3

@@ -12,14 +12,18 @@
     module PG
       module Patches
         # Module to prepend to PG::Connection for instrumentation
         module Connection # rubocop:disable Metrics/ModuleLength
           PG::Constants::EXEC_ISH_METHODS.each do |method|
-            define_method method do |*args|
+            define_method method do |*args, &block|
               span_name, attrs = span_attrs(:query, *args)
               tracer.in_span(span_name, attributes: attrs, kind: :client) do
-                super(*args)
+                if block
+                  block.call(super(*args))
+                else
+                  super(*args)
+                end
               end
             end
           end
 
           PG::Constants::PREPARE_ISH_METHODS.each do |method|
@@ -30,13 +34,17 @@
               end
             end
           end
 
           PG::Constants::EXEC_PREPARED_ISH_METHODS.each do |method|
-            define_method method do |*args|
+            define_method method do |*args, &block|
               span_name, attrs = span_attrs(:execute, *args)
               tracer.in_span(span_name, attributes: attrs, kind: :client) do
-                super(*args)
+                if block
+                  block.call(super(*args))
+                else
+                  super(*args)
+                end
               end
             end
           end
 
           private