ddtrace 0.54.2 → 1.8.0

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

@@ -0,0 +1,33 @@
+# typed: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Pg
+        # pg integration constants
+        # @public_api Changing resource names, tag names, or environment variables creates breaking changes.
+        module Ext
+          ENV_ENABLED = 'DD_TRACE_PG_ENABLED'.freeze
+          ENV_SERVICE_NAME = 'DD_TRACE_PG_SERVICE_NAME'.freeze
+          ENV_ANALYTICS_ENABLED = 'DD_TRACE_PG_ANALYTICS_ENABLED'.freeze
+          ENV_ANALYTICS_SAMPLE_RATE = 'DD_TRACE_PG_ANALYTICS_SAMPLE_RATE'.freeze
+          DEFAULT_PEER_SERVICE_NAME = 'pg'.freeze
+          SPAN_EXEC = 'pg.exec'.freeze
+          SPAN_EXEC_PARAMS = 'pg.exec.params'.freeze
+          SPAN_EXEC_PREPARED = 'pg.exec.prepared'.freeze
+          SPAN_ASYNC_EXEC = 'pg.async.exec'.freeze
+          SPAN_ASYNC_EXEC_PARAMS = 'pg.async.exec.params'.freeze
+          SPAN_ASYNC_EXEC_PREPARED = 'pg.async.exec.prepared'.freeze
+          SPAN_SYNC_EXEC = 'pg.sync.exec'.freeze
+          SPAN_SYNC_EXEC_PARAMS = 'pg.sync.exec.params'.freeze
+          SPAN_SYNC_EXEC_PREPARED = 'pg.sync.exec.prepared'.freeze
+          TAG_DB_NAME = 'pg.db.name'.freeze
+          TAG_COMPONENT = 'pg'.freeze
+          TAG_OPERATION_QUERY = 'query'.freeze
+
+          TAG_SYSTEM = 'postgresql'.freeze
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/pg/instrumentation.rb

@@ -0,0 +1,155 @@
+# typed: false
+
+require_relative '../../metadata/ext'
+require_relative '../analytics'
+require_relative '../ext'
+require_relative 'ext'
+
+require_relative '../propagation/sql_comment'
+require_relative '../propagation/sql_comment/mode'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Pg
+        # PG::Connection patch module
+        module Instrumentation
+          def self.included(base)
+            base.prepend(InstanceMethods)
+          end
+
+          # PG::Connection patch methods
+          module InstanceMethods
+            def exec(sql, *args)
+              trace(Ext::SPAN_EXEC, sql: sql) do |sql_statement|
+                super(sql_statement, *args)
+              end
+            end
+
+            def exec_params(sql, params, *args)
+              trace(Ext::SPAN_EXEC_PARAMS, sql: sql) do |sql_statement|
+                super(sql_statement, params, *args)
+              end
+            end
+
+            def exec_prepared(statement_name, params, *args)
+              trace(Ext::SPAN_EXEC_PREPARED, statement_name: statement_name) do
+                super(statement_name, params, *args)
+              end
+            end
+
+            def async_exec(sql, *args)
+              trace(Ext::SPAN_ASYNC_EXEC, sql: sql) do |sql_statement|
+                super(sql_statement, *args)
+              end
+            end
+
+            def async_exec_params(sql, params, *args)
+              trace(Ext::SPAN_ASYNC_EXEC_PARAMS, sql: sql) do |sql_statement|
+                super(sql_statement, params, *args)
+              end
+            end
+
+            def async_exec_prepared(statement_name, params, *args)
+              trace(Ext::SPAN_ASYNC_EXEC_PREPARED, statement_name: statement_name) do
+                super(statement_name, params, *args)
+              end
+            end
+
+            def sync_exec(sql, *args)
+              trace(Ext::SPAN_SYNC_EXEC, sql: sql) do |sql_statement|
+                super(sql_statement, *args)
+              end
+            end
+
+            def sync_exec_params(sql, params, *args)
+              trace(Ext::SPAN_SYNC_EXEC_PARAMS, sql: sql) do |sql_statement|
+                super(sql_statement, params, *args)
+              end
+            end
+
+            def sync_exec_prepared(statement_name, params, *args)
+              trace(Ext::SPAN_SYNC_EXEC_PREPARED, statement_name: statement_name) do
+                super(statement_name, params, *args)
+              end
+            end
+
+            private
+
+            def trace(name, sql: nil, statement_name: nil)
+              service = Datadog.configuration_for(self, :service_name) || datadog_configuration[:service_name]
+              resource = statement_name || sql
+
+              Tracing.trace(
+                name,
+                service: service,
+                resource: resource,
+                type: Tracing::Metadata::Ext::SQL::TYPE
+              ) do |span, trace_op|
+                annotate_span_with_query!(span, service)
+                # Set analytics sample rate
+                Contrib::Analytics.set_sample_rate(span, analytics_sample_rate) if analytics_enabled?
+
+                if sql
+                  propagation_mode = Contrib::Propagation::SqlComment::Mode.new(comment_propagation)
+                  Contrib::Propagation::SqlComment.annotate!(span, propagation_mode)
+                  propagated_sql_statement = Contrib::Propagation::SqlComment.prepend_comment(
+                    sql,
+                    span,
+                    trace_op,
+                    propagation_mode
+                  )
+                end
+
+                result = yield(propagated_sql_statement)
+                annotate_span_with_result!(span, result)
+                result
+              end
+            end
+
+            def annotate_span_with_query!(span, service)
+              span.set_tag(Ext::TAG_DB_NAME, db)
+
+              span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
+              span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_QUERY)
+              span.set_tag(Tracing::Metadata::Ext::TAG_KIND, Tracing::Metadata::Ext::SpanKind::TAG_CLIENT)
+
+              # Tag as an external peer service
+              span.set_tag(Tracing::Metadata::Ext::TAG_PEER_SERVICE, service)
+              span.set_tag(Tracing::Metadata::Ext::TAG_PEER_HOSTNAME, host)
+
+              span.set_tag(Contrib::Ext::DB::TAG_INSTANCE, db)
+              span.set_tag(Contrib::Ext::DB::TAG_USER, user)
+              span.set_tag(Contrib::Ext::DB::TAG_SYSTEM, Ext::TAG_SYSTEM)
+
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_HOST, host)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_PORT, port)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_DESTINATION_NAME, host)
+              span.set_tag(Tracing::Metadata::Ext::NET::TAG_DESTINATION_PORT, port)
+            end
+
+            def annotate_span_with_result!(span, result)
+              span.set_tag(Contrib::Ext::DB::TAG_ROW_COUNT, result.ntuples)
+            end
+
+            def datadog_configuration
+              Datadog.configuration.tracing[:pg]
+            end
+
+            def analytics_enabled?
+              datadog_configuration[:analytics_enabled]
+            end
+
+            def analytics_sample_rate
+              datadog_configuration[:analytics_sample_rate]
+            end
+
+            def comment_propagation
+              datadog_configuration[:comment_propagation]
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,43 @@
+# typed: false
+
+require_relative '../integration'
+require_relative 'configuration/settings'
+require_relative 'patcher'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Pg
+        # Description of pg integration
+        class Integration
+          include Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new('0.18.4')
+
+          # @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

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

@@ -0,0 +1,31 @@
+# typed: ignore
+
+require_relative '../patcher'
+require_relative 'instrumentation'
+
+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
+            patch_pg_connection
+          end
+
+          def patch_pg_connection
+            ::PG::Connection.include(Instrumentation)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/presto/configuration/settings.rb

@@ -0,0 +1,38 @@
+# typed: false
+
+require_relative '../../configuration/settings'
+require_relative '../ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Presto
+        module Configuration
+          # Custom settings for the Presto integration
+          # @public_api
+          class Settings < Contrib::Configuration::Settings
+            option :enabled do |o|
+              o.default { env_to_bool(Ext::ENV_ENABLED, true) }
+              o.lazy
+            end
+
+            option :analytics_enabled do |o|
+              o.default { env_to_bool(Ext::ENV_ANALYTICS_ENABLED, false) }
+              o.lazy
+            end
+
+            option :analytics_sample_rate do |o|
+              o.default { env_to_float(Ext::ENV_ANALYTICS_SAMPLE_RATE, 1.0) }
+              o.lazy
+            end
+
+            option :service_name do |o|
+              o.default { ENV.fetch(Ext::ENV_SERVICE_NAME, Ext::DEFAULT_PEER_SERVICE_NAME) }
+              o.lazy
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/presto/ext.rb

@@ -0,0 +1,33 @@
+# typed: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Presto
+        # Presto integration constants
+        # @public_api Changing resource names, tag names, or environment variables creates breaking changes.
+        module Ext
+          ENV_ENABLED = 'DD_TRACE_PRESTO_ENABLED'.freeze
+          ENV_SERVICE_NAME = 'DD_TRACE_PRESTO_SERVICE_NAME'.freeze
+          ENV_ANALYTICS_ENABLED = 'DD_TRACE_PRESTO_ANALYTICS_ENABLED'.freeze
+          ENV_ANALYTICS_SAMPLE_RATE = 'DD_TRACE_PRESTO_ANALYTICS_SAMPLE_RATE'.freeze
+          DEFAULT_PEER_SERVICE_NAME = 'presto'.freeze
+          SPAN_QUERY = 'presto.query'.freeze
+          SPAN_KILL = 'presto.kill_query'.freeze
+          TAG_SCHEMA_NAME = 'presto.schema'.freeze
+          TAG_CATALOG_NAME = 'presto.catalog'.freeze
+          TAG_USER_NAME = 'presto.user'.freeze
+          TAG_TIME_ZONE = 'presto.time_zone'.freeze
+          TAG_LANGUAGE = 'presto.language'.freeze
+          TAG_PROXY = 'presto.http_proxy'.freeze
+          TAG_MODEL_VERSION = 'presto.model_version'.freeze
+          TAG_QUERY_ID = 'presto.query.id'.freeze
+          TAG_QUERY_ASYNC = 'presto.query.async'.freeze
+          TAG_COMPONENT = 'presto'.freeze
+          TAG_OPERATION_QUERY = 'query'.freeze
+          TAG_OPERATION_KILL = 'kill'.freeze
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/presto/instrumentation.rb

@@ -0,0 +1,124 @@
+# typed: true
+
+require_relative '../../metadata/ext'
+require_relative 'ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Presto
+        # Instrumentation for Presto integration
+        module Instrumentation
+          # Instrumentation for Presto::Client::Client
+          module Client
+            def self.included(base)
+              base.prepend(InstanceMethods)
+            end
+
+            # Instance methods for Presto::Client
+            module InstanceMethods
+              def run(query)
+                Tracing.trace(
+                  Ext::SPAN_QUERY,
+                  service: datadog_configuration[:service_name]
+                ) do |span|
+                  begin
+                    decorate!(span, Ext::TAG_OPERATION_QUERY)
+                    span.resource = query
+                    span.span_type = Tracing::Metadata::Ext::SQL::TYPE
+                    span.set_tag(Ext::TAG_QUERY_ASYNC, false)
+                  rescue StandardError => e
+                    Datadog.logger.debug("error preparing span for presto: #{e}")
+                  end
+
+                  super(query)
+                end
+              end
+
+              def query(query, &blk)
+                Tracing.trace(
+                  Ext::SPAN_QUERY,
+                  service: datadog_configuration[:service_name]
+                ) do |span|
+                  begin
+                    decorate!(span, Ext::TAG_OPERATION_QUERY)
+                    span.resource = query
+                    span.span_type = Tracing::Metadata::Ext::SQL::TYPE
+                    span.set_tag(Ext::TAG_QUERY_ASYNC, !blk.nil?)
+                  rescue StandardError => e
+                    Datadog.logger.debug("error preparing span for presto: #{e}")
+                  end
+
+                  super(query, &blk)
+                end
+              end
+
+              def kill(query_id)
+                Tracing.trace(
+                  Ext::SPAN_KILL,
+                  service: datadog_configuration[:service_name]
+                ) do |span|
+                  begin
+                    decorate!(span, Ext::TAG_OPERATION_KILL)
+                    span.resource = Ext::SPAN_KILL
+                    span.span_type = Tracing::Metadata::Ext::AppTypes::TYPE_DB
+                    # ^ not an SQL type span, since there's no SQL query
+                    span.set_tag(Ext::TAG_QUERY_ID, query_id)
+                  rescue StandardError => e
+                    Datadog.logger.debug("error preparing span for presto: #{e}")
+                  end
+
+                  super(query_id)
+                end
+              end
+
+              private
+
+              def datadog_configuration
+                Datadog.configuration.tracing[:presto]
+              end
+
+              def decorate!(span, operation)
+                span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
+                span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, operation)
+
+                if (host_port = @options[:server])
+                  host, port = Core::Utils.extract_host_port(host_port)
+                  if host && port
+                    span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_HOST, host)
+                    span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_PORT, port)
+
+                    span.set_tag(Tracing::Metadata::Ext::TAG_PEER_HOSTNAME, host)
+                  else
+                    span.set_tag(Tracing::Metadata::Ext::NET::TAG_TARGET_HOST, host_port)
+                    span.set_tag(Tracing::Metadata::Ext::TAG_PEER_HOSTNAME, host_port)
+                  end
+                end
+
+                set_nilable_tag!(span, :user, Ext::TAG_USER_NAME)
+                set_nilable_tag!(span, :schema, Ext::TAG_SCHEMA_NAME)
+                set_nilable_tag!(span, :catalog, Ext::TAG_CATALOG_NAME)
+                set_nilable_tag!(span, :time_zone, Ext::TAG_TIME_ZONE)
+                set_nilable_tag!(span, :language, Ext::TAG_LANGUAGE)
+                set_nilable_tag!(span, :http_proxy, Ext::TAG_PROXY)
+                set_nilable_tag!(span, :model_version, Ext::TAG_MODEL_VERSION)
+
+                # Tag as an external peer service
+                span.set_tag(Tracing::Metadata::Ext::TAG_PEER_SERVICE, span.service)
+
+                # Set analytics sample rate
+                if Contrib::Analytics.enabled?(datadog_configuration[:analytics_enabled])
+                  Contrib::Analytics.set_sample_rate(span, datadog_configuration[:analytics_sample_rate])
+                end
+              end
+
+              def set_nilable_tag!(span, key, tag_name)
+                @options[key].tap { |val| span.set_tag(tag_name, val) if val }
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/presto/integration.rb

@@ -0,0 +1,43 @@
+# typed: false
+
+require_relative '../integration'
+require_relative 'configuration/settings'
+require_relative 'patcher'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Presto
+        # Description of Presto integration
+        class Integration
+          include Contrib::Integration
+
+          MINIMUM_VERSION = Gem::Version.new('0.5.14')
+
+          # @public_api Changing the integration name or integration options can cause breaking changes
+          register_as :presto
+
+          def self.version
+            Gem.loaded_specs['presto-client'] && Gem.loaded_specs['presto-client'].version
+          end
+
+          def self.loaded?
+            !defined?(::Presto::Client::Client).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

data/lib/datadog/tracing/contrib/presto/patcher.rb

@@ -0,0 +1,37 @@
+# typed: true
+
+require_relative '../../../core/utils/only_once'
+require_relative '../patcher'
+require_relative 'ext'
+require_relative 'instrumentation'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Presto
+        # Patcher enables patching of 'presto-client' module.
+        module Patcher
+          include Contrib::Patcher
+
+          PATCH_ONLY_ONCE = Core::Utils::OnlyOnce.new
+
+          module_function
+
+          def patched?
+            PATCH_ONLY_ONCE.ran?
+          end
+
+          def patch
+            PATCH_ONLY_ONCE.run do
+              begin
+                ::Presto::Client::Client.include(Instrumentation::Client)
+              rescue StandardError => e
+                Datadog.logger.error("Unable to apply Presto integration: #{e}")
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/propagation/sql_comment/comment.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+# typed: false
+
+require 'erb'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Propagation
+        module SqlComment
+          # To be prepended to a sql statement.
+          class Comment
+            def initialize(hash)
+              @hash = hash
+            end
+
+            def to_s
+              @string ||= begin
+                ret = String.new
+
+                @hash.each do |key, value|
+                  next if value.nil?
+
+                  # Url encode
+                  value = ERB::Util.url_encode(value)
+
+                  # Escape SQL
+                  ret << "#{key}='#{value}',"
+                end
+
+                # Remove the last `,`
+                ret.chop!
+
+                "/*#{ret}*/"
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/propagation/sql_comment/ext.rb

@@ -0,0 +1,33 @@
+# typed: false
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Propagation
+        module SqlComment
+          module Ext
+            ENV_DBM_PROPAGATION_MODE = 'DD_DBM_PROPAGATION_MODE'.freeze
+
+            # The default mode for sql comment propagation
+            DISABLED = 'disabled'.freeze
+
+            # The `service` mode propagates service configuration
+            SERVICE = 'service'.freeze
+
+            # The `full` mode propagates service configuration + trace context
+            FULL = 'full'.freeze
+
+            # The value should be `true` when `full` mode
+            TAG_DBM_TRACE_INJECTED = '_dd.dbm_trace_injected'.freeze
+
+            KEY_DATABASE_SERVICE = 'dddbs'.freeze
+            KEY_ENVIRONMENT = 'dde'.freeze
+            KEY_PARENT_SERVICE = 'ddps'.freeze
+            KEY_VERSION = 'ddpv'.freeze
+            KEY_TRACEPARENT = 'traceparent'.freeze
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/propagation/sql_comment/mode.rb

@@ -0,0 +1,28 @@
+# typed: false
+
+require_relative 'ext'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Propagation
+        # Implements sql comment propagation related contracts.
+        module SqlComment
+          Mode = Struct.new(:mode) do
+            def enabled?
+              service? || full?
+            end
+
+            def service?
+              mode == Ext::SERVICE
+            end
+
+            def full?
+              mode == Ext::FULL
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/contrib/propagation/sql_comment.rb

@@ -0,0 +1,47 @@
+# typed: false
+
+require_relative 'sql_comment/comment'
+require_relative 'sql_comment/ext'
+
+require_relative '../../distributed/trace_context'
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Propagation
+        # Implements sql comment propagation related contracts.
+        module SqlComment
+          def self.annotate!(span_op, mode)
+            return unless mode.enabled?
+
+            span_op.set_tag(Ext::TAG_DBM_TRACE_INJECTED, true) if mode.full?
+          end
+
+          # Inject span_op and trace_op instead of TraceDigest to improve memory usage
+          # for `disabled` and `service` mode
+          def self.prepend_comment(sql, span_op, trace_op, mode)
+            return sql unless mode.enabled?
+
+            tags = {
+              Ext::KEY_DATABASE_SERVICE => span_op.service,
+              Ext::KEY_ENVIRONMENT => datadog_configuration.env,
+              Ext::KEY_PARENT_SERVICE => datadog_configuration.service,
+              Ext::KEY_VERSION => datadog_configuration.version
+            }
+
+            if mode.full?
+              tags[Ext::KEY_TRACEPARENT] =
+                Tracing::Distributed::TraceContext.new(fetcher: nil).send(:build_traceparent, trace_op.to_digest)
+            end
+
+            "#{Comment.new(tags)} #{sql}"
+          end
+
+          def self.datadog_configuration
+            Datadog.configuration
+          end
+        end
+      end
+    end
+  end
+end