Initial implementation of Presto intergration

Author: ahammel

.env

@@ -20,5 +20,7 @@ TEST_POSTGRES_HOST=127.0.0.1
 TEST_POSTGRES_PASSWORD=postgres
 TEST_POSTGRES_PORT=5432
 TEST_POSTGRES_USER=postgres
+TEST_PRESTO_HOST=localhost
+TEST_PRESTO_PORT=8080
 TEST_REDIS_HOST=127.0.0.1
 TEST_REDIS_PORT=6379

Appraisals

@@ -406,6 +406,7 @@ elsif Gem::Version.new('2.2.0') <= Gem::Version.new(RUBY_VERSION) \
       gem 'hiredis'
       gem 'mongo', '>= 2.8.0'
       gem 'mysql2', '< 0.5', platform: :ruby
+      gem 'presto-client', '>=  0.5.14'
       gem 'racecar', '>= 0.3.5'
       gem 'rack'
       gem 'rack-test'
@@ -545,6 +546,7 @@ elsif Gem::Version.new('2.3.0') <= Gem::Version.new(RUBY_VERSION) \
       gem 'hiredis'
       gem 'mongo', '>= 2.8.0'
       gem 'mysql2', '< 0.5', platform: :ruby
+      gem 'presto-client', '>=  0.5.14'
       gem 'racecar', '>= 0.3.5'
       gem 'rack'
       gem 'rack-test'
@@ -578,6 +580,7 @@ elsif Gem::Version.new('2.4.0') <= Gem::Version.new(RUBY_VERSION)
       gem 'hiredis'
       gem 'mongo', '>= 2.8.0'
       gem 'mysql2', '< 0.5', platform: :ruby
+      gem 'presto-client', '>=  0.5.14'
       gem 'racecar', '>= 0.3.5'
       gem 'rack'
       gem 'rack-test'

Rakefile

@@ -75,6 +75,7 @@ namespace :spec do
     :http,
     :mongodb,
     :mysql2,
+    :presto,
     :racecar,
     :rack,
     :rake,

docker-compose.yml

@@ -150,6 +150,7 @@ services:
       - mongodb
       - mysql
       - postgres
+      - presto
       - redis
     env_file: ./.env
     environment:
@@ -161,6 +162,8 @@ services:
       - TEST_MONGODB_HOST=mongodb
       - TEST_MYSQL_HOST=mysql
       - TEST_POSTGRES_HOST=postgres
+      - TEST_PRESTO_HOST=presto
+      - TEST_PRESTO_PORT=8080
       - TEST_REDIS_HOST=redis
     stdin_open: true
     tty: true
@@ -225,6 +228,12 @@ services:
       - "5432"
     ports:
       - "${TEST_POSTGRES_PORT}:5432"
+  presto:
+    image: starburstdata/presto:302-e.2 
+    expose:
+      - "8080"
+    ports:
+      - "${TEST_PRESTO_PORT}:8080"
   redis:
     image: redis:3.0
     expose:

lib/ddtrace.rb

@@ -35,6 +35,7 @@ module Datadog
 require 'ddtrace/contrib/grpc/integration'
 require 'ddtrace/contrib/http/integration'
 require 'ddtrace/contrib/integration'
+require 'ddtrace/contrib/presto/integration'
 require 'ddtrace/contrib/mysql2/integration'
 require 'ddtrace/contrib/mongodb/integration'
 require 'ddtrace/contrib/racecar/integration'

lib/ddtrace/contrib/presto/configuration/settings.rb

@@ -0,0 +1,23 @@
+require 'ddtrace/contrib/configuration/settings'
+require 'ddtrace/contrib/presto/ext'
+
+module Datadog
+  module Contrib
+    module Presto
+      module Configuration
+        # Custom settings for the Presto integration
+        class Settings < Contrib::Configuration::Settings
+          option  :analytics_enabled,
+                  default: -> { env_to_bool(Ext::ENV_ANALYTICS_ENABLED, false) },
+                  lazy: true
+
+          option  :analytics_sample_rate,
+                  default: -> { env_to_float(Ext::ENV_ANALYTICS_SAMPLE_RATE, 1.0) },
+                  lazy: true
+
+          option :service_name, default: Ext::SERVICE_NAME
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/presto/ext.rb

@@ -0,0 +1,23 @@
+module Datadog
+  module Contrib
+    module Presto
+      # MongoDB integration constants
+      module Ext
+        APP = 'presto'.freeze
+        ENV_ANALYTICS_ENABLED = 'DD_PRESTO_ANALYTICS_ENABLED'.freeze
+        ENV_ANALYTICS_SAMPLE_RATE = 'DD_PRESTO_ANALYTICS_SAMPLE_RATE'.freeze
+        SERVICE_NAME = 'presto'.freeze
+        SPAN_QUERY = 'presto.query'.freeze
+        SPAN_RESUME = 'presto.resume_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
+      end
+    end
+  end
+end

lib/ddtrace/contrib/presto/instrumentation.rb

@@ -0,0 +1,70 @@
+require 'ddtrace/pin'
+require 'ddtrace/ext/net'
+require 'ddtrace/ext/sql'
+require 'ddtrace/ext/app_types'
+require 'ddtrace/contrib/presto/ext'
+
+module Datadog
+  module Contrib
+    module Presto
+      # Instrumentation for Presto integration
+      module Instrumentation
+        # Instrumentation for Presto::Client::Client
+        module Client
+          def self.included(base)
+            base.send(:prepend, InstanceMethods)
+          end
+
+          # Instance methods for Presto::Client
+          module InstanceMethods
+            def run(query)
+              datadog_pin.tracer.trace(Ext::SPAN_QUERY) do |span|
+                decorate!(span)
+                super(query)
+              end
+            end
+
+            def query(query, &blk)
+              datadog_pin.tracer.trace(Ext::SPAN_QUERY) do |span|
+                decorate!(span)
+                super(query, &blk)
+              end
+            end
+
+            def kill(query_id)
+              datadog_pin.tracer.trace(Ext::SPAN_KILL) do |span|
+                decorate!(span)
+                super(query_id)
+              end
+            end
+
+            private
+
+            def decorate!(span)
+              span.service = datadog_pin.service
+              span.span_type = Datadog::Ext::SQL::TYPE
+
+              span.set_tag(Ext::TAG_SCHEMA_NAME, @options[:schema])
+              span.set_tag(Ext::TAG_CATALOG_NAME, @options[:catalog])
+              span.set_tag(Ext::TAG_USER_NAME, @options[:user])
+              span.set_tag(Ext::TAG_TIME_ZONE, @options[:time_zone])
+              span.set_tag(Ext::TAG_LANGUAGE, @options[:language])
+              span.set_tag(Ext::TAG_PROXY, @options[:http_proxy])
+              span.set_tag(Ext::TAG_MODEL_VERSION, @options[:model_version])
+              span.set_tag(Datadog::Ext::NET::TARGET_HOST, @options[:server])
+            end
+
+            def datadog_pin
+              @datadog_pin ||= Datadog::Pin.new(
+                Datadog.configuration[:presto][:service_name],
+                app: Ext::APP,
+                app_type: Datadog::Ext::AppTypes::DB,
+                tracer: Datadog.configuration[:presto][:tracer]
+              )
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/presto/integration.rb

@@ -0,0 +1,32 @@
+require 'ddtrace/contrib/integration'
+require 'ddtrace/contrib/presto/configuration/settings'
+require 'ddtrace/contrib/presto/patcher'
+
+module Datadog
+  module Contrib
+    module Presto
+      # Description of Presto integration
+      class Integration
+        include Contrib::Integration
+
+        register_as :presto
+
+        def self.version
+          Gem.loaded_specs['presto-client'] && Gem.loaded_specs['presto-client'].version
+        end
+
+        def self.present?
+          super && defined?(::Presto::Client::Client)
+        end
+
+        def default_configuration
+          Configuration::Settings.new
+        end
+
+        def patcher
+          Patcher
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/presto/patcher.rb

@@ -0,0 +1,30 @@
+require 'ddtrace/contrib/patcher'
+require 'ddtrace/contrib/presto/ext'
+require 'ddtrace/contrib/presto/instrumentation'
+
+module Datadog
+  module Contrib
+    module Presto
+      # Patcher enables patching of 'presto-client' module.
+      module Patcher
+        include Contrib::Patcher
+
+        module_function
+
+        def patched?
+          done?(:presto)
+        end
+
+        def patch
+          do_once(:presto) do
+            begin
+              ::Presto::Client::Client.send(:include, Instrumentation::Client)
+            rescue StandardError => e
+              Datadog::Tracer.log.error("Unable to apply Presto integration: #{e}")
+            end
+          end
+        end
+      end
+    end
+  end
+end