feat: ActiveRecord instrumentation improvements

instrumentation/active_record/README.md

@@ -30,6 +30,42 @@ OpenTelemetry::SDK.configure do |c|
 end
 ```
 
+## Configuration Options
+
+The instrumentation supports the following configuration options:
+
+- **enable_notifications_instrumentation:** Enables instrumentation of SQL queries using ActiveSupport notifications. When enabled, generates spans for each SQL query with additional metadata including operation names, async status, and caching information.
+  - Default: `false`
+
+## Active Support Instrumentation
+
+This instrumentation can optionally leverage `ActiveSupport::Notifications` to provide detailed SQL query instrumentation. When enabled via the `enable_notifications_instrumentation` configuration option, it subscribes to `sql.active_record` events to create spans for individual SQL queries.
+
+### Enabling SQL Notifications
+
+```ruby
+OpenTelemetry::SDK.configure do |c|
+  c.use 'OpenTelemetry::Instrumentation::ActiveRecord',
+        enable_notifications_instrumentation: true
+end
+```
+
+See the table below for details of what [Rails ActiveRecord Events](https://guides.rubyonrails.org/active_support_instrumentation.html#active-record) are recorded by this instrumentation:
+
+| Event Name | Creates Span? | Notes |
+| - | - | - |
+| `sql.active_record` | :white_check_mark: | Creates an `internal` span for each SQL query with operation name, async status, and caching information |
+
+### SQL Query Spans
+
+When notifications instrumentation is enabled, each SQL query executed through ActiveRecord generates a span with:
+
+- **Span name**: Derived from the query operation (e.g., `"User Create"`, `"Account Load"`, `"Post Update"`)
+- **Span kind**: `internal`
+- **Attributes**:
+  - `rails.active_record.query.async` (boolean): Present and set to `true` for asynchronous queries
+  - `rails.active_record.query.cached` (boolean): Present and set to `true` for cached query results
+
 ## Examples
 
 Example usage can be seen in the [`./example/trace_demonstration.rb` file](https://github.com/open-telemetry/opentelemetry-ruby-contrib/blob/main/instrumentation/active_record/example/trace_demonstration.rb)

instrumentation/active_record/lib/opentelemetry/instrumentation/active_record/handlers.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require_relative 'handlers/sql_handler'
+
+module OpenTelemetry
+  module Instrumentation
+    module ActiveRecord
+      # Module that contains custom event handlers for ActiveRecord notifications
+      module Handlers
+        module_function
+
+        # Subscribes Event Handlers to relevant ActiveRecord notifications
+        #
+        # The following events are recorded as spans:
+        # - sql.active_record
+        #
+        # @note this method is not thread safe and should not be used in a multi-threaded context
+        def subscribe
+          return unless Array(@subscriptions).empty?
+
+          config = ActiveRecord::Instrumentation.instance.config
+          return unless config[:enable_notifications_instrumentation]
+
+          sql_handler = Handlers::SqlHandler.new
+
+          @subscriptions = [
+            ::ActiveSupport::Notifications.subscribe('sql.active_record', sql_handler)
+          ]
+        end
+
+        # Removes Event Handler Subscriptions for ActiveRecord notifications
+        # @note this method is not thread-safe and should not be used in a multi-threaded context
+        def unsubscribe
+          @subscriptions&.each { |subscriber| ::ActiveSupport::Notifications.unsubscribe(subscriber) }
+          @subscriptions = nil
+        end
+      end
+    end
+  end
+end

instrumentation/active_record/lib/opentelemetry/instrumentation/active_record/handlers/sql_handler.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+module OpenTelemetry
+  module Instrumentation
+    module ActiveRecord
+      module Handlers
+        # SqlHandler handles sql.active_record ActiveSupport notifications
+        class SqlHandler
+          # Invoked by ActiveSupport::Notifications at the start of the instrumentation block
+          #
+          # @param name [String] of the Event
+          # @param id [String] of the event
+          # @param payload [Hash] containing SQL execution information
+          # @return [Hash] the payload passed as a method argument
+          def start(name, id, payload)
+            span = tracer.start_span(
+              name,
+              kind: :internal,
+              attributes: { 'db.operation' => payload[:name] || 'SQL',
+                            'rails.active_record.query.async' => payload[:async] == true,
+                            'rails.active_record.query.cached' => payload[:cached] == true }
+            )
+            token = OpenTelemetry::Context.attach(
+              OpenTelemetry::Trace.context_with_span(span)
+            )
+            payload[:__opentelemetry_span] = span
+            payload[:__opentelemetry_ctx_token] = token
+          rescue StandardError => e
+            OpenTelemetry.handle_error(exception: e)
+          end
+
+          # Invoked by ActiveSupport::Notifications at the end of the instrumentation block
+          #
+          # @param name [String] of the Event
+          # @param id [String] of the event
+          # @param payload [Hash] containing SQL execution information
+          def finish(name, id, payload)
+            span = payload.delete(:__opentelemetry_span)
+            token = payload.delete(:__opentelemetry_ctx_token)
+            return unless span && token
+
+            if (e = payload[:exception_object])
+              span.record_exception(e)
+              span.status = OpenTelemetry::Trace::Status.error('Unhandled exception')
+            end
+
+            span.finish
+            OpenTelemetry::Context.detach(token)
+          rescue StandardError => e
+            OpenTelemetry.handle_error(exception: e)
+          end
+
+          private
+
+          def tracer
+            OpenTelemetry::Instrumentation::ActiveRecord::Instrumentation.instance.tracer
+          end
+        end
+      end
+    end
+  end
+end

instrumentation/active_record/lib/opentelemetry/instrumentation/active_record/instrumentation.rb

@@ -11,9 +11,10 @@ module ActiveRecord
       class Instrumentation < OpenTelemetry::Instrumentation::Base
         MINIMUM_VERSION = Gem::Version.new('7')
 
-        install do |_config|
+        install do |config|
           require_dependencies
           patch_activerecord
+          subscribe_to_notifications if config[:enable_notifications_instrumentation]
         end
 
         present do
@@ -24,6 +25,8 @@ class Instrumentation < OpenTelemetry::Instrumentation::Base
           gem_version >= MINIMUM_VERSION
         end
 
+        option :enable_notifications_instrumentation, default: false, validate: :boolean
+
         private
 
         def gem_version
@@ -39,6 +42,7 @@ def require_dependencies
           require_relative 'patches/transactions_class_methods'
           require_relative 'patches/validations'
           require_relative 'patches/relation_persistence'
+          require_relative 'handlers'
         end
 
         def patch_activerecord
@@ -57,6 +61,10 @@ def patch_activerecord
             ::ActiveRecord::Relation.prepend(Patches::RelationPersistence)
           end
         end
+
+        def subscribe_to_notifications
+          Handlers.subscribe
+        end
       end
     end
   end

instrumentation/active_record/test/instrumentation/active_record/handlers/sql_handler_test.rb

@@ -0,0 +1,189 @@
+# frozen_string_literal: true
+
+# Copyright The OpenTelemetry Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+require 'test_helper'
+require 'opentelemetry/instrumentation/active_record/handlers'
+
+describe OpenTelemetry::Instrumentation::ActiveRecord::Handlers::SqlHandler do
+  let(:instrumentation) { OpenTelemetry::Instrumentation::ActiveRecord::Instrumentation.instance }
+  let(:config) { { enable_notifications_instrumentation: true } }
+  let(:exporter) { EXPORTER }
+  let(:spans) { exporter.finished_spans.select { |s| s.name == 'sql.active_record' } }
+
+  before do
+    # Capture original config before modification
+    @original_config = instrumentation.config.dup
+
+    OpenTelemetry::Instrumentation::ActiveRecord::Handlers.unsubscribe
+    instrumentation.instance_variable_set(:@config, config)
+    instrumentation.instance_variable_set(:@installed, false)
+
+    instrumentation.install(config)
+    User.delete_all
+    Account.delete_all
+    exporter.reset
+  end
+
+  after do
+    # Restore original configuration and reinstall
+    OpenTelemetry::Instrumentation::ActiveRecord::Handlers.unsubscribe
+    instrumentation.instance_variable_set(:@config, @original_config)
+    instrumentation.instance_variable_set(:@installed, false)
+    instrumentation.install(@original_config)
+
+    # Clear any test data
+    User.delete_all
+    Account.delete_all
+  end
+
+  describe 'when sql.active_record notifications are emitted' do
+    it 'creates spans with operation name from payload' do
+      User.create!(name: 'otel')
+
+      _(spans).wont_be_empty
+    end
+
+    it 'records async attribute when query is async' do
+      # Create a user first so there's data to load
+      Account.transaction do
+        account = Account.create!
+        User.create!(name: 'otel', account: account)
+      end
+
+      exporter.reset
+
+      ActiveRecord::Base.asynchronous_queries_tracker.start_session
+      relations = [
+        User.limit(1).load_async,
+        User.where(name: 'otel').includes(:account).load_async
+      ]
+      # Now wait for completion
+      result = relations.flat_map(&:to_a)
+      ActiveRecord::Base.asynchronous_queries_tracker.finalize_session(true)
+
+      _(result).wont_be_empty
+
+      # Trigger a real async query
+
+      # The query should have run asynchronously
+      async_spans = spans.select { |span| span.attributes['rails.active_record.query.async'] == true }
+      _(async_spans).wont_be_empty
+    end
+
+    it 'records cached attribute when query is cached' do
+      # First query - not cached
+      User.first
+
+      exporter.reset
+
+      # Second query with caching enabled - should be cached
+      User.cache do
+        User.first
+        User.first
+      end
+
+      cached_spans = spans.select { |s| s.attributes['rails.active_record.query.cached'] == true }
+      _(cached_spans).wont_be_empty
+    end
+
+    it 'records synchronous queries' do
+      _(User.all.to_a).must_be_empty
+
+      values = spans.map { |span| span.attributes['rails.active_record.query.async'] }.uniq
+      _(values).must_equal [false]
+    end
+
+    it 'records actual queries' do
+      _(User.all.to_a).must_be_empty
+
+      values = spans.map { |span| span.attributes['rails.active_record.query.cached'] }.uniq
+      _(values).must_equal [false]
+    end
+
+    it 'records exceptions on spans' do
+      # Create a scenario that will cause a SQL error
+      begin
+        ActiveRecord::Base.connection.execute('SELECT * FROM nonexistent_table')
+      rescue StandardError
+        # Expected to fail
+      end
+
+      error_spans = spans.select { |s| s.status.code == OpenTelemetry::Trace::Status::ERROR }
+      _(error_spans).wont_be_empty
+    end
+
+    it 'sets span kind to internal' do
+      User.first
+
+      sql_spans = spans.reject { |s| s.attributes['db.operation'] == 'ActiveRecord::Base.transaction' }
+      _(sql_spans).wont_be_empty
+
+      sql_spans.each do |span|
+        _(span.kind).must_equal :internal
+      end
+    end
+
+    it 'uses SQL as default name when name is not present' do
+      # Manually trigger a notification without a name
+      ActiveRecord::Base.connection.execute('SELECT 1')
+
+      _(spans.map { |s| s.attributes['db.operation'] }).must_equal ['SQL']
+    end
+
+    it 'creates nested spans correctly' do
+      Account.transaction do
+        account = Account.create!
+        User.create!(name: 'otel', account: account)
+      end
+
+      # Verify parent-child relationships
+      transaction_span = spans.find { |s| s.attributes['db.operation'] == 'TRANSACTION' }
+      _(transaction_span).wont_be_nil
+
+      create_spans = spans.select { |s| s.attributes['db.operation'].include?('Create') }
+
+      _(create_spans.map { |s| s.attributes['db.operation'] }).must_equal(['Account Create', 'User Create'])
+      _(create_spans.map(&:parent_span_id)).must_equal([transaction_span.span_id, transaction_span.span_id])
+    end
+  end
+
+  describe 'with complex queries' do
+    before do
+      Account.create!
+      5.times { User.create!(name: 'otel') }
+    end
+
+    it 'instruments SELECT queries' do
+      User.where(name: 'otel').first
+
+      select_spans = spans.select { |s| s.attributes['db.operation'].include?('User Load') }
+      _(select_spans).wont_be_empty
+    end
+
+    it 'instruments UPDATE queries' do
+      user = User.first
+      user.update!(counter: 42)
+
+      update_spans = spans.select { |s| s.attributes['db.operation'].include?('User Update') }
+      _(update_spans).wont_be_empty
+    end
+
+    it 'instruments DELETE queries' do
+      user = User.first
+      user.destroy
+
+      delete_spans = spans.select { |s| s.attributes['db.operation'].include?('User Destroy') }
+      _(delete_spans).wont_be_empty
+    end
+
+    it 'instruments batch operations' do
+      User.where(name: 'otel').delete_all
+
+      delete_spans = spans.select { |s| s.attributes['db.operation'].include?('SQL') || s.attributes['db.operation'].include?('Delete') }
+      _(delete_spans).wont_be_empty
+    end
+  end
+end