Added Sequel support for execute_ddl, execute_dui and execute_insert

lib/ddtrace/contrib/sequel/patcher.rb

@@ -1,4 +1,5 @@
 module Datadog
+  # rubocop:disable Metrics/ModuleLength
   module Contrib
     module Sequel
       # Patcher enables patching of 'sequel' module.
@@ -35,6 +36,7 @@ def patch
           @patched
         end
 
+        # rubocop:disable Metrics/MethodLength
         def patch_sequel_database
           ::Sequel::Database.send(:include, Datadog::Contrib::Sequel::Utils)
           ::Sequel::Database.class_eval do
@@ -50,9 +52,11 @@ def initialize(*args)
             end
 
             alias_method :run_without_datadog, :run
+            remove_method :run
+
             def run(sql, options = ::Sequel::OPTS)
               pin = Datadog::Pin.get_from(self)
-              return run_without_datadog(req, body) unless pin && pin.tracer
+              return run_without_datadog(sql, options) unless pin && pin.tracer
 
               opts = parse_opts(sql, options)
 
@@ -70,7 +74,9 @@ def run(sql, options = ::Sequel::OPTS)
           end
         end
 
+        # rubocop:disable Metrics/AbcSize
         def patch_sequel_dataset
+          # rubocop:disable Metrics/BlockLength
           ::Sequel::Dataset.send(:include, Datadog::Contrib::Sequel::Utils)
           ::Sequel::Dataset.class_eval do
             alias_method :initialize_without_datadog, :initialize
@@ -85,9 +91,11 @@ def initialize(*args)
             end
 
             alias_method :execute_without_datadog, :execute
+            remove_method :execute
+
             def execute(sql, options = ::Sequel::OPTS, &block)
               pin = Datadog::Pin.get_from(self)
-              return execute_without_datadog(req, body, &block) unless pin && pin.tracer
+              return execute_without_datadog(sql, options, &block) unless pin && pin.tracer
 
               opts = parse_opts(sql, options)
               response = nil
@@ -101,6 +109,66 @@ def execute(sql, options = ::Sequel::OPTS, &block)
               end
               response
             end
+
+            alias_method :execute_ddl_without_datadog, :execute_ddl
+            remove_method :execute_ddl
+
+            def execute_ddl(sql, options = ::Sequel::OPTS, &block)
+              pin = Datadog::Pin.get_from(self)
+              return execute_ddl_without_datadog(sql, options, &block) unless pin && pin.tracer
+
+              opts = parse_opts(sql, options)
+              response = nil
+
+              pin.tracer.trace('sequel.query') do |span|
+                span.service = pin.service
+                span.resource = opts[:query]
+                span.span_type = Datadog::Ext::SQL::TYPE
+                span.set_tag('sequel.db.vendor', adapter_name)
+                response = execute_ddl_without_datadog(sql, options, &block)
+              end
+              response
+            end
+
+            alias_method :execute_dui_without_datadog, :execute_dui
+            remove_method :execute_dui
+
+            def execute_dui(sql, options = ::Sequel::OPTS, &block)
+              pin = Datadog::Pin.get_from(self)
+              return execute_dui_without_datadog(sql, options, &block) unless pin && pin.tracer
+
+              opts = parse_opts(sql, options)
+              response = nil
+
+              pin.tracer.trace('sequel.query') do |span|
+                span.service = pin.service
+                span.resource = opts[:query]
+                span.span_type = Datadog::Ext::SQL::TYPE
+                span.set_tag('sequel.db.vendor', adapter_name)
+                response = execute_dui_without_datadog(sql, options, &block)
+              end
+              response
+            end
+
+            alias_method :execute_insert_without_datadog, :execute_insert
+            remove_method :execute_insert
+
+            def execute_insert(sql, options = ::Sequel::OPTS, &block)
+              pin = Datadog::Pin.get_from(self)
+              return execute_insert_without_datadog(sql, options, &block) unless pin && pin.tracer
+
+              opts = parse_opts(sql, options)
+              response = nil
+
+              pin.tracer.trace('sequel.query') do |span|
+                span.service = pin.service
+                span.resource = opts[:query]
+                span.span_type = Datadog::Ext::SQL::TYPE
+                span.set_tag('sequel.db.vendor', adapter_name)
+                response = execute_insert_without_datadog(sql, options, &block)
+              end
+              response
+            end
           end
         end
       end

lib/ddtrace/contrib/sequel/utils.rb

@@ -1,3 +1,5 @@
+require 'ddtrace/contrib/rails/utils'
+
 module Datadog
   module Contrib
     module Sequel

test/contrib/sequel/miniapp_test.rb

@@ -19,9 +19,13 @@ def check_span_process(span, parent_id, trace_id)
     assert_equal(trace_id, span.trace_id)
   end
 
-  def check_span_command(span, parent_id, trace_id)
+  def check_span_command(span, parent_id, trace_id, resource)
     assert_equal('sequel.query', span.name)
     assert_equal('sequel', span.service)
+    assert_equal('sql', span.span_type)
+    assert_equal('sqlite', span.get_tag('sequel.db.vendor'))
+    assert_equal(resource, span.resource)
+    assert_equal(0, span.status)
     assert_equal(parent_id, span.parent_id)
     assert_equal(trace_id, span.trace_id)
   end
@@ -42,19 +46,28 @@ def test_miniapp
       tracer.trace('process') do |subspan|
         subspan.service = 'datalayer'
         subspan.resource = 'home'
-        sequel[:table].insert(name: 'data')
+        sequel[:table].insert(name: 'data1')
+        sequel[:table].insert(name: 'data2')
+        data = sequel[:table].select.to_a
+        assert_equal(2, data.length)
+        data.each do |row|
+          assert_match(/^data.$/, row[:name])
+        end
       end
     end
 
     spans = tracer.writer.spans
 
-    assert_equal(3, spans.length)
-    process, publish, sequel_cmd = spans
+    assert_equal(6, spans.length)
+    process, publish, sequel_cmd1, sequel_cmd2, sequel_cmd3, sequel_cmd4 = spans
     check_span_publish publish
     parent_id = publish.span_id
     trace_id = publish.trace_id
     check_span_process process, parent_id, trace_id
     parent_id = process.span_id
-    check_span_command sequel_cmd, parent_id, trace_id
+    check_span_command sequel_cmd1, parent_id, trace_id, 'INSERT INTO `table` (`name`) VALUES (?)'
+    check_span_command sequel_cmd2, parent_id, trace_id, 'INSERT INTO `table` (`name`) VALUES (?)'
+    check_span_command sequel_cmd3, parent_id, trace_id, 'SELECT * FROM `table`'
+    check_span_command sequel_cmd4, parent_id, trace_id, 'SELECT sqlite_version()'
   end
 end