Add Braze PG instrumentation

lib/datadog/tracing/contrib.rb

@@ -57,6 +57,7 @@ module Contrib
 require 'datadog/tracing/contrib/lograge/integration'
 require 'datadog/tracing/contrib/mongodb/integration'
 require 'datadog/tracing/contrib/mysql2/integration'
+require 'datadog/tracing/contrib/pg/integration'
 require 'datadog/tracing/contrib/presto/integration'
 require 'datadog/tracing/contrib/qless/integration'
 require 'datadog/tracing/contrib/que/integration'

lib/datadog/tracing/contrib/pg/configuration/settings.rb

@@ -0,0 +1,19 @@
+# typed: false
+
+require 'datadog/tracing/contrib/configuration/settings'
+require 'datadog/tracing/contrib/pg/ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Pg
+        module Configuration
+          # Custom settings for the PG integration
+          class Settings < Contrib::Configuration::Settings
+            option :service_name, default: Ext::SERVICE_NAME
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/pg/connection.rb

@@ -0,0 +1,88 @@
+# typed: false
+
+require 'datadog/tracing/metadata/ext'
+require 'datadog/tracing/contrib/pg/ext'
+
+module Datadog
+  module Contrib
+    module Pg
+      # PG::Connection patch module
+      module Connection
+        module_function
+
+        def included(base)
+          base.send(:prepend, InstanceMethods)
+        end
+
+        # PG::Connection patch instance methods
+        module InstanceMethods
+
+          # sync_exec(sql) -> PG::Result
+          # sync_exec(sql) {|pg_result| block}
+          def sync_exec(sql)
+            service = Datadog.configuration_for(self, :service_name) || datadog_configuration[:service_name]
+            Tracing.trace(Ext::SPAN_QUERY, service: service) do |span|
+              span.resource = sql
+              span.span_type = Tracing::Metadata::Ext::SQL::TYPE
+              span.set_tag(Ext::TAG_DB_NAME, db)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_HOST, host)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_PORT, port)
+              super # this will pass all args, including the block
+            end
+          end
+
+          # sync_exec_params(sql, params[, result_format[, type_map]] ) -> PG::Result
+          # sync_exec_params(sql, params[, result_format[, type_map]] ) {|pg_result| block }
+          # exec_params (and async version) is parsed with rb_scan_args like so:
+          # rb_scan_args(argc, argv, "22", &command, &paramsData.params, &in_res_fmt, &paramsData.typemap);
+          # meaning it expects 2 required arguments and 2 explicit optional
+          def sync_exec_params(sql, params, result_format = nil, type_map = nil)
+            service = Datadog.configuration_for(self, :service_name) || datadog_configuration[:service_name]
+            Tracing.trace(Ext::SPAN_QUERY, service: service) do |span|
+              span.resource = sql
+              span.span_type = Tracing::Metadata::Ext::SQL::TYPE
+              span.set_tag(Ext::TAG_DB_NAME, db)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_HOST, host)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_PORT, port)
+              super # this will pass all args, including the block
+            end
+          end
+
+          # async_exec(sql) -> PG::Result OR
+          # async_exec(sql) {|pg_result| block}
+          def async_exec(sql)
+            service = Datadog.configuration_for(self, :service_name) || datadog_configuration[:service_name]
+            Tracing.trace(Ext::SPAN_QUERY, service: service) do |span|
+              span.resource = sql
+              span.span_type = Tracing::Metadata::Ext::SQL::TYPE
+              span.set_tag(Ext::TAG_DB_NAME, db)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_HOST, host)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_PORT, port)
+              super # this will pass all args, including the block
+            end
+          end
+
+          # async_exec_params(sql, params[, result_format[, type_map]] ) -> PG::Result
+          # async_exec_params(sql, params[, result_format[, type_map]] ) {|pg_result| block }
+          def async_exec_params(sql, params, result_format = nil, type_map = nil)
+            service = Datadog.configuration_for(self, :service_name) || datadog_configuration[:service_name]
+            Tracing.trace(Ext::SPAN_QUERY, service: service) do |span|
+              span.resource = sql
+              span.span_type = Tracing::Metadata::Ext::SQL::TYPE
+              span.set_tag(Ext::TAG_DB_NAME, db)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_HOST, host)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_PORT, port)
+              super # this will pass all args, including the block
+            end
+          end
+
+          private
+
+          def datadog_configuration
+            Datadog.configuration.tracing[:pg]
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/pg/ext.rb

@@ -0,0 +1,17 @@
+# typed: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Pg
+        # PG integration constants
+        module Ext
+          APP = 'pg'.freeze
+          SERVICE_NAME = 'pg'.freeze
+          SPAN_QUERY = 'pg.query'.freeze
+          TAG_DB_NAME = 'pg.db.name'.freeze
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/pg/integration.rb

@@ -0,0 +1,43 @@
+# typed: false
+
+require 'datadog/tracing/contrib/integration'
+require 'datadog/tracing/contrib/pg/configuration/settings'
+require 'datadog/tracing/contrib/pg/patcher'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Pg
+        # Description of PG integration
+        class Integration
+          include Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new('1.3.5')
+
+          # @public_api Changing the integration name or integration options can cause breaking changes
+          register_as :pg
+
+          def self.version
+            Gem.loaded_specs['pg'] && Gem.loaded_specs['pg'].version
+          end
+
+          def self.loaded?
+            !defined?(::Pg).nil?
+          end
+
+          def self.compatible?
+            super && version >= MINIMUM_VERSION
+          end
+
+          def new_configuration
+            Configuration::Settings.new
+          end
+
+          def patcher
+            Patcher
+          end
+        end
+      end
+    end
+  end
+end

lib/datadog/tracing/contrib/pg/patcher.rb

@@ -0,0 +1,27 @@
+# typed: false
+
+require 'datadog/tracing/contrib/patcher'
+require 'datadog/tracing/contrib/pg/connection'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Pg
+        # Patcher enables patching of 'pg' module.
+        module Patcher
+          include Contrib::Patcher
+
+          module_function
+
+          def target_version
+            Integration.version
+          end
+
+          def patch
+            ::PG::Connection.send(:include, Connection)
+          end
+        end
+      end
+    end
+  end
+end