Add ActiveSupport::Notifications::Event helper

lib/ddtrace/contrib/active_record/event.rb

@@ -0,0 +1,30 @@
+require 'ddtrace/contrib/active_support/notifications/event'
+
+module Datadog
+  module Contrib
+    module ActiveRecord
+      # Defines basic behaviors for an ActiveRecord event.
+      module Event
+        def self.included(base)
+          base.send(:include, ActiveSupport::Notifications::Event)
+          base.send(:extend, ClassMethods)
+        end
+
+        # Class methods for ActiveRecord events.
+        module ClassMethods
+          def span_options
+            { service: configuration[:service_name] }
+          end
+
+          def tracer
+            configuration[:tracer]
+          end
+
+          def configuration
+            Datadog.configuration[:active_record]
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/active_record/events.rb

@@ -0,0 +1,30 @@
+require 'ddtrace/contrib/active_record/events/instantiation'
+require 'ddtrace/contrib/active_record/events/sql'
+
+module Datadog
+  module Contrib
+    module ActiveRecord
+      # Defines collection of instrumented ActiveRecord events
+      module Events
+        ALL = [
+          Events::Instantiation,
+          Events::SQL
+        ].freeze
+
+        module_function
+
+        def all
+          self::ALL
+        end
+
+        def subscriptions
+          all.collect(&:subscriptions).collect(&:to_a).flatten
+        end
+
+        def subscribe!
+          all.each(&:subscribe!)
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/active_record/events/instantiation.rb

@@ -0,0 +1,51 @@
+require 'ddtrace/contrib/active_record/event'
+
+module Datadog
+  module Contrib
+    module ActiveRecord
+      module Events
+        # Defines instrumentation for instantiation.active_record event
+        module Instantiation
+          include ActiveRecord::Event
+
+          EVENT_NAME = 'instantiation.active_record'.freeze
+          SPAN_NAME = 'active_record.instantiation'.freeze
+          DEFAULT_SERVICE_NAME = 'active_record'.freeze
+
+          module_function
+
+          def supported?
+            Gem.loaded_specs['activerecord'] \
+              && Gem.loaded_specs['activerecord'].version >= Gem::Version.new('4.2')
+          end
+
+          def event_name
+            self::EVENT_NAME
+          end
+
+          def span_name
+            self::SPAN_NAME
+          end
+
+          def process(span, event, _id, payload)
+            # Inherit service name from parent, if available.
+            span.service = if configuration[:orm_service_name]
+                             configuration[:orm_service_name]
+                           elsif span.parent
+                             span.parent.service
+                           else
+                             self::DEFAULT_SERVICE_NAME
+                           end
+
+            span.resource = payload.fetch(:class_name)
+            span.span_type = 'custom'
+            span.set_tag('active_record.instantiation.class_name', payload.fetch(:class_name))
+            span.set_tag('active_record.instantiation.record_count', payload.fetch(:record_count))
+          rescue StandardError => e
+            Datadog::Tracer.log.debug(e.message)
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/active_record/events/sql.rb

@@ -0,0 +1,48 @@
+require 'ddtrace/contrib/active_record/event'
+
+module Datadog
+  module Contrib
+    module ActiveRecord
+      module Events
+        # Defines instrumentation for sql.active_record event
+        module SQL
+          include ActiveRecord::Event
+
+          EVENT_NAME = 'sql.active_record'.freeze
+          SPAN_NAME = 'active_record.sql'.freeze
+
+          module_function
+
+          def event_name
+            self::EVENT_NAME
+          end
+
+          def span_name
+            self::SPAN_NAME
+          end
+
+          def process(span, event, _id, payload)
+            connection_config = Utils.connection_config(payload[:connection_id])
+            span.name = "#{connection_config[:adapter_name]}.query"
+            span.service = configuration[:service_name]
+            span.resource = payload.fetch(:sql)
+            span.span_type = Datadog::Ext::SQL::TYPE
+
+            # Find out if the SQL query has been cached in this request. This meta is really
+            # helpful to users because some spans may have 0ns of duration because the query
+            # is simply cached from memory, so the notification is fired with start == finish.
+            cached = payload[:cached] || (payload[:name] == 'CACHE')
+
+            span.set_tag('active_record.db.vendor', connection_config[:adapter_name])
+            span.set_tag('active_record.db.name', connection_config[:database_name])
+            span.set_tag('active_record.db.cached', cached) if cached
+            span.set_tag('out.host', connection_config[:adapter_host])
+            span.set_tag('out.port', connection_config[:adapter_port])
+          rescue StandardError => e
+            Datadog::Tracer.log.debug(e.message)
+          end
+        end
+      end
+    end
+  end
+end

lib/ddtrace/contrib/active_record/patcher.rb

@@ -1,18 +1,14 @@
 require 'ddtrace/ext/sql'
 require 'ddtrace/ext/app_types'
 require 'ddtrace/contrib/active_record/utils'
-require 'ddtrace/contrib/active_support/notifications/subscriber'
+require 'ddtrace/contrib/active_record/events'
 
 module Datadog
   module Contrib
     module ActiveRecord
       # Patcher enables patching of 'active_record' module.
       module Patcher
         include Base
-        include ActiveSupport::Notifications::Subscriber
-
-        NAME_SQL = 'sql.active_record'.freeze
-        NAME_INSTANTIATION = 'instantiation.active_record'.freeze
 
         register_as :active_record, auto_patch: false
         option :service_name, depends_on: [:tracer] do |value|
@@ -24,36 +20,14 @@ module Patcher
         option :tracer, default: Datadog.tracer do |value|
           (value || Datadog.tracer).tap do |v|
             # Make sure to update tracers of all subscriptions
-            subscriptions.each do |subscription|
+            Events.subscriptions.each do |subscription|
               subscription.tracer = v
             end
           end
         end
 
         @patched = false
 
-        on_subscribe do
-          # sql.active_record
-          subscribe(
-            self::NAME_SQL,                         # Event name
-            'active_record.sql',                    # Span name
-            { service: get_option(:service_name) }, # Span options
-            get_option(:tracer),                    # Tracer
-            &method(:sql)                           # Handler
-          )
-
-          # instantiation.active_record
-          if instantiation_tracing_supported?
-            subscribe(
-              self::NAME_INSTANTIATION,               # Event name
-              'active_record.instantiation',          # Span name
-              { service: get_option(:service_name) }, # Span options
-              get_option(:tracer),                    # Tracer
-              &method(:instantiation)                 # Handler
-            )
-          end
-        end
-
         module_function
 
         # patched? tells whether patch has been successfully applied
@@ -64,7 +38,7 @@ def patched?
         def patch
           if !@patched && defined?(::ActiveRecord)
             begin
-              subscribe!
+              Events.subscribe!
               @patched = true
             rescue StandardError => e
               Datadog::Tracer.log.error("Unable to apply Active Record integration: #{e}")
@@ -73,50 +47,6 @@ def patch
 
           @patched
         end
-
-        def instantiation_tracing_supported?
-          Gem.loaded_specs['activerecord'] \
-            && Gem.loaded_specs['activerecord'].version >= Gem::Version.new('4.2')
-        end
-
-        def sql(span, event, _id, payload)
-          connection_config = Utils.connection_config(payload[:connection_id])
-          span.name = "#{connection_config[:adapter_name]}.query"
-          span.service = get_option(:service_name)
-          span.resource = payload.fetch(:sql)
-          span.span_type = Datadog::Ext::SQL::TYPE
-
-          # Find out if the SQL query has been cached in this request. This meta is really
-          # helpful to users because some spans may have 0ns of duration because the query
-          # is simply cached from memory, so the notification is fired with start == finish.
-          cached = payload[:cached] || (payload[:name] == 'CACHE')
-
-          span.set_tag('active_record.db.vendor', connection_config[:adapter_name])
-          span.set_tag('active_record.db.name', connection_config[:database_name])
-          span.set_tag('active_record.db.cached', cached) if cached
-          span.set_tag('out.host', connection_config[:adapter_host])
-          span.set_tag('out.port', connection_config[:adapter_port])
-        rescue StandardError => e
-          Datadog::Tracer.log.debug(e.message)
-        end
-
-        def instantiation(span, event, _id, payload)
-          # Inherit service name from parent, if available.
-          span.service = if get_option(:orm_service_name)
-                           get_option(:orm_service_name)
-                         elsif span.parent
-                           span.parent.service
-                         else
-                           'active_record'
-                         end
-
-          span.resource = payload.fetch(:class_name)
-          span.span_type = 'custom'
-          span.set_tag('active_record.instantiation.class_name', payload.fetch(:class_name))
-          span.set_tag('active_record.instantiation.record_count', payload.fetch(:record_count))
-        rescue StandardError => e
-          Datadog::Tracer.log.debug(e.message)
-        end
       end
     end
   end

lib/ddtrace/contrib/active_support/notifications/event.rb

@@ -0,0 +1,62 @@
+require 'ddtrace/contrib/active_support/notifications/subscriber'
+
+module Datadog
+  module Contrib
+    module ActiveSupport
+      module Notifications
+        # Defines behaviors for an ActiveSupport::Notifications event.
+        # Compose this into a module or class, then define
+        # #event_name, #span_name, and #process. You can then
+        # invoke Event.subscribe! to more easily subscribe to an event.
+        module Event
+          def self.included(base)
+            base.send(:include, Subscriber)
+            base.send(:extend, ClassMethods)
+            base.send(:on_subscribe) { base.subscribe }
+          end
+
+          # Redefines some class behaviors for a Subscriber to make
+          # it a bit simpler for an Event.
+          module ClassMethods
+            def subscribe!
+              super
+            end
+
+            def subscription(span_name = nil, options = nil, tracer = nil)
+              super(
+                span_name || self.span_name,
+                options || span_options,
+                tracer || self.tracer,
+                &method(:process)
+              )
+            end
+
+            def subscribe(pattern = nil, span_name = nil, options = nil, tracer = nil)
+              if supported?
+                super(
+                  pattern || event_name,
+                  span_name || self.span_name,
+                  options || span_options,
+                  tracer || self.tracer,
+                  &method(:process)
+                )
+              end
+            end
+
+            def supported?
+              true
+            end
+
+            def span_options
+              {}
+            end
+
+            def tracer
+              Datadog.tracer
+            end
+          end
+        end
+      end
+    end
+  end
+end