rails-otel-context 0.5.0

data/lib/rails_otel_context/activerecord_context.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+module RailsOtelContext
+  # Extracts ActiveRecord model name, method, and scope from sql.active_record
+  # notifications and scope instrumentation.
+  #
+  # Two mechanisms work together:
+  # 1. sql.active_record subscriber — captures model name + AR operation type
+  #    (e.g., "Transaction Load") at query time via payload[:name].
+  # 2. Scope tracking — wraps scope-generated methods to store the scope name
+  #    on the Relation object, then captures it at materialization time via
+  #    Relation#exec_queries. This handles lazy scopes like
+  #    Transaction.recent_completed.to_a where the scope method returns before
+  #    SQL fires.
+  module ActiveRecordContext
+    THREAD_KEY       = :_rails_otel_ctx_ar
+    SCOPE_THREAD_KEY = :_rails_otel_ctx_scope
+    private_constant :THREAD_KEY, :SCOPE_THREAD_KEY
+
+    # Subscriber for sql.active_record notifications.
+    class Subscriber
+      def start(_name, _id, payload)
+        ar_name = payload[:name]
+        return unless ar_name
+        return if ar_name == 'SCHEMA' || ar_name.start_with?('CACHE') || ar_name == 'SQL'
+
+        ctx = ActiveRecordContext.parse_ar_name(ar_name)
+        return unless ctx
+
+        # Include scope name if one was captured by RelationScopeCapture
+        scope = Thread.current[SCOPE_THREAD_KEY]
+        ctx[:scope_name] = scope if scope
+        Thread.current[THREAD_KEY] = ctx
+      end
+
+      def finish(_name, _id, _payload)
+        Thread.current[THREAD_KEY] = nil
+      end
+    end
+
+    # Wraps scope-generated class methods to store the scope name on the Relation.
+    module ScopeNameTracking
+      def scope(name, body, &)
+        super
+
+        # Guard against double-wrapping on class reload in development
+        @_otel_wrapped_scopes ||= {}
+        return if @_otel_wrapped_scopes[name]
+
+        name_str = name.to_s.freeze
+        original = method(name)
+        define_singleton_method(name) do |*args|
+          relation = original.call(*args)
+          if relation.is_a?(::ActiveRecord::Relation)
+            relation.instance_variable_set(:@_otel_scope_name, name_str)
+          end
+          relation
+        end
+        @_otel_wrapped_scopes[name] = true
+      end
+    end
+
+    # Captures scope name from Relation at SQL materialization time.
+    module RelationScopeCapture
+      def exec_queries(&)
+        scope_name = instance_variable_get(:@_otel_scope_name)
+        Thread.current[SCOPE_THREAD_KEY] = scope_name if scope_name
+        super
+      ensure
+        Thread.current[SCOPE_THREAD_KEY] = nil
+      end
+    end
+
+    module_function
+
+    def install!
+      return unless defined?(::ActiveSupport::Notifications)
+      return unless defined?(::ActiveRecord::Base)
+
+      ActiveSupport::Notifications.subscribe('sql.active_record', Subscriber.new)
+      ::ActiveRecord::Base.extend(ScopeNameTracking)
+      ::ActiveRecord::Relation.prepend(RelationScopeCapture)
+    end
+
+    def current
+      Thread.current[THREAD_KEY]
+    end
+
+    def clear!
+      Thread.current[THREAD_KEY] = nil
+      Thread.current[SCOPE_THREAD_KEY] = nil
+    end
+
+    # Test helpers: set AR context directly for unit tests.
+    def stub_context(context)
+      Thread.current[THREAD_KEY] = context
+    end
+
+    def stub_scope(scope_name)
+      Thread.current[SCOPE_THREAD_KEY] = scope_name
+    end
+    private :stub_scope
+
+    # Parses "Transaction Load" → { model_name: "Transaction", method_name: "Load" }
+    def parse_ar_name(name)
+      return nil unless name
+
+      parts = name.split(' ', 2)
+      return nil unless parts.size == 2
+
+      model_name = parts[0]
+      method_name = parts[1]
+
+      return nil if model_name == 'ActiveRecord'
+
+      { model_name: model_name, method_name: method_name }
+    end
+
+    # Legacy method for adapter tests.
+    def extract(app_root: nil) # rubocop:disable Lint/UnusedMethodArgument
+      current
+    end
+  end
+end

data/lib/rails_otel_context/adapters/clickhouse.rb

@@ -0,0 +1,133 @@
+# frozen_string_literal: true
+
+module RailsOtelContext
+  module Adapters
+    module Clickhouse
+      module_function
+
+      CANDIDATE_METHODS = %i[query select insert execute command].freeze
+
+      def install!(app_root:, threshold_ms:)
+        begin
+          require 'click_house'
+        rescue LoadError
+          # ClickHouse client gem is optional for consumers.
+        end
+
+        target_clients.each do |klass|
+          methods = CANDIDATE_METHODS.select { |method_name| klass.method_defined?(method_name) }
+          next if methods.empty?
+
+          patch_module = patch_module_for(klass, methods)
+          patch_module.configure(app_root: app_root, threshold_ms: threshold_ms)
+          next if klass.ancestors.include?(patch_module)
+
+          klass.prepend(patch_module)
+        end
+      end
+
+      def target_clients
+        clients = []
+
+        clients << ::ClickHouse::Client if defined?(::ClickHouse::Client)
+        clients << ::ClickHouse::Connection if defined?(::ClickHouse::Connection)
+        clients << ::Clickhouse::Client if defined?(::Clickhouse::Client)
+
+        clients.compact.uniq
+      end
+
+      def patch_module_for(klass, methods)
+        @patch_modules ||= {}
+        key = [klass.name, methods.sort].join(':')
+        @patch_modules[key] ||= build_patch_module(methods)
+      end
+
+      def build_patch_module(methods)
+        mod = Module.new do
+          class << self
+            attr_accessor :app_root, :threshold_ms
+
+            def configure(app_root:, threshold_ms:)
+              @app_root = app_root.to_s
+              @threshold_ms = threshold_ms.to_f
+            end
+
+            def source_location_for_app
+              return unless Thread.respond_to?(:each_caller_location)
+
+              Thread.each_caller_location do |location|
+                path = location.absolute_path || location.path
+                next unless path&.start_with?(app_root)
+                next if path.include?('/gems/')
+
+                return [path.delete_prefix("#{app_root}/"), location.lineno]
+              end
+
+              nil
+            end
+
+            def activerecord_context
+              RailsOtelContext::ActiveRecordContext.extract(app_root: app_root)
+            end
+          end
+
+          methods.each do |method_name|
+            define_method(method_name) do |*args, &block|
+              if Thread.current[:rails_otel_context_clickhouse_instrumenting]
+                return super(*args, &block)
+              end
+
+              source = mod.source_location_for_app
+              ar_context = mod.activerecord_context
+              started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+              operation = method_name.to_s.upcase
+              statement = args.first.is_a?(String) ? args.first : nil
+
+              tracer = OpenTelemetry.tracer_provider.tracer('rails-otel-context-clickhouse')
+              Thread.current[:rails_otel_context_clickhouse_instrumenting] = true
+
+              tracer.in_span("#{operation} clickhouse", kind: :client) do |span|
+                span.set_attribute('db.system', 'clickhouse')
+                span.set_attribute('db.operation', operation)
+                span.set_attribute('db.statement', statement) if statement
+
+                # Rename span if formatter is configured and AR context is available
+                if ar_context && RailsOtelContext.configuration.span_name_formatter
+                  begin
+                    new_name = RailsOtelContext.configuration.span_name_formatter.call(span.name, ar_context)
+                    span.name = new_name if new_name && new_name != span.name && span.respond_to?(:name=)
+                  rescue StandardError => e
+                    warn "[RailsOtelContext] Span name formatter error: #{e.message}"
+                  end
+                end
+
+                result = super(*args, &block)
+                duration_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at) * 1000.0
+
+                if source && duration_ms >= mod.threshold_ms
+                  span.set_attribute('code.filepath', source[0])
+                  span.set_attribute('code.lineno', source[1])
+                  span.set_attribute('db.query.duration_ms', duration_ms.round(1))
+                  span.set_attribute('db.query.slow_threshold_ms', mod.threshold_ms)
+                end
+
+                # Add ActiveRecord context if available
+                if ar_context
+                  span.set_attribute('code.activerecord.model', ar_context[:model_name]) if ar_context[:model_name]
+                  span.set_attribute('code.activerecord.method', ar_context[:method_name]) if ar_context[:method_name]
+                end
+
+                result
+              end
+            ensure
+              Thread.current[:rails_otel_context_clickhouse_instrumenting] = false
+            end
+          end
+        end
+
+        mod
+      end
+      private_class_method :build_patch_module
+    end
+  end
+end

data/lib/rails_otel_context/adapters/mysql2.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module RailsOtelContext
+  module Adapters
+    module Mysql2
+      module_function
+
+      def install!(app_root:, threshold_ms:)
+        return unless defined?(::Mysql2::Client)
+
+        patch_module = patch_module_for
+        patch_module.configure(app_root: app_root, threshold_ms: threshold_ms)
+
+        return if ::Mysql2::Client.ancestors.include?(patch_module)
+
+        ::Mysql2::Client.prepend(patch_module)
+      end
+
+      def patch_module_for
+        @patch_module ||= build_patch_module
+      end
+
+      def build_patch_module
+        mod = Module.new do
+          class << self
+            include RailsOtelContext::SourceLocation
+
+            attr_accessor :app_root, :threshold_ms
+
+            def configure(app_root:, threshold_ms:)
+              @app_root = app_root.to_s
+              @threshold_ms = threshold_ms.to_f
+            end
+          end
+
+          # AR context and span renaming handled by CallContextProcessor.apply_db_context.
+          # Adapters only handle slow query source location tracking.
+          define_method(:query) do |sql, options = {}|
+            source = mod.source_location_for_app
+            started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            result = super(sql, options)
+            duration_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at) * 1000.0
+
+            span = OpenTelemetry::Trace.current_span
+            if span.context.valid?
+              threshold = RailsOtelContext.configuration.mysql2_slow_query_threshold_ms
+              if source && duration_ms >= threshold
+                span.set_attribute('code.filepath', source[0])
+                span.set_attribute('code.lineno', source[1])
+                span.set_attribute('db.query.duration_ms', duration_ms.round(1))
+                span.set_attribute('db.query.slow_threshold_ms', threshold)
+              end
+            end
+
+            result
+          end
+
+          define_method(:prepare) do |sql|
+            source = mod.source_location_for_app
+            started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            result = super(sql)
+            duration_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at) * 1000.0
+
+            span = OpenTelemetry::Trace.current_span
+            if span.context.valid?
+              threshold = RailsOtelContext.configuration.mysql2_slow_query_threshold_ms
+              if source && duration_ms >= threshold
+                span.set_attribute('code.filepath', source[0])
+                span.set_attribute('code.lineno', source[1])
+                span.set_attribute('db.query.duration_ms', duration_ms.round(1))
+                span.set_attribute('db.query.slow_threshold_ms', threshold)
+              end
+            end
+
+            result
+          end
+        end
+
+        mod
+      end
+      private_class_method :build_patch_module
+    end
+  end
+end

data/lib/rails_otel_context/adapters/pg.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module RailsOtelContext
+  module Adapters
+    module PG
+      module_function
+
+      def install!(app_root:, threshold_ms:)
+        return unless defined?(::PG::Connection)
+
+        methods = exec_methods
+        return if methods.empty?
+
+        patch_module = patch_module_for(methods)
+        patch_module.configure(app_root: app_root, threshold_ms: threshold_ms)
+
+        return if ::PG::Connection.ancestors.include?(patch_module)
+
+        ::PG::Connection.prepend(patch_module)
+      end
+
+      def exec_methods
+        return [] unless defined?(::PG::Constants::EXEC_ISH_METHODS)
+        return [] unless defined?(::PG::Constants::EXEC_PREPARED_ISH_METHODS)
+
+        (::PG::Constants::EXEC_ISH_METHODS + ::PG::Constants::EXEC_PREPARED_ISH_METHODS).uniq
+      end
+
+      def patch_module_for(methods)
+        @patch_module ||= build_patch_module(methods)
+      end
+
+      def build_patch_module(methods)
+        mod = Module.new do
+          class << self
+            include RailsOtelContext::SourceLocation
+
+            attr_accessor :app_root, :threshold_ms
+
+            def configure(app_root:, threshold_ms:)
+              @app_root = app_root.to_s
+              @threshold_ms = threshold_ms.to_f
+            end
+          end
+
+          # AR context and span renaming handled by CallContextProcessor.apply_db_context.
+          # Adapters only handle slow query source location tracking.
+          methods.each do |method_name|
+            define_method(method_name) do |*args, &user_block|
+              source = mod.source_location_for_app
+              started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+              super(*args) do |result|
+                duration_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at) * 1000.0
+
+                span = OpenTelemetry::Trace.current_span
+                threshold = RailsOtelContext.configuration.pg_slow_query_threshold_ms
+                if source && duration_ms >= threshold
+                  span.set_attribute('code.filepath', source[0])
+                  span.set_attribute('code.lineno', source[1])
+                  span.set_attribute('db.query.duration_ms', duration_ms.round(1))
+                  span.set_attribute('db.query.slow_threshold_ms', threshold)
+                end
+
+                user_block ? user_block.call(result) : result
+              end
+            end
+          end
+        end
+
+        mod
+      end
+      private_class_method :build_patch_module
+    end
+  end
+end

data/lib/rails_otel_context/adapters/redis.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module RailsOtelContext
+  module Adapters
+    module Redis
+      module_function
+
+      def install!(app_root:)
+        return unless defined?(::RedisClient::Middlewares)
+        return unless defined?(::OpenTelemetry::Instrumentation::Redis)
+
+        patch_module = patch_module_for
+        patch_module.configure(app_root: app_root)
+
+        return if ::RedisClient::Middlewares.ancestors.include?(patch_module)
+
+        ::RedisClient::Middlewares.prepend(patch_module)
+      end
+
+      def patch_module_for
+        @patch_module ||= build_patch_module
+      end
+
+      def build_patch_module
+        mod = Module.new do
+          class << self
+            attr_accessor :app_root
+
+            def configure(app_root:)
+              @app_root = app_root.to_s
+            end
+
+            def source_location_for_app
+              return unless Thread.respond_to?(:each_caller_location)
+
+              Thread.each_caller_location do |location|
+                path = location.absolute_path || location.path
+                next unless path&.start_with?(app_root)
+                next if path.include?('/gems/')
+
+                return [path.delete_prefix("#{app_root}/"), location.lineno]
+              end
+
+              nil
+            end
+          end
+
+          define_method(:call) do |command, redis_config, &block|
+            source = mod.source_location_for_app
+            return super(command, redis_config, &block) unless source
+
+            OpenTelemetry::Instrumentation::Redis.with_attributes(
+              'code.filepath' => source[0],
+              'code.lineno' => source[1]
+            ) do
+              super(command, redis_config, &block)
+            end
+          end
+
+          define_method(:call_pipelined) do |commands, redis_config, &block|
+            source = mod.source_location_for_app
+            return super(commands, redis_config, &block) unless source
+
+            OpenTelemetry::Instrumentation::Redis.with_attributes(
+              'code.filepath' => source[0],
+              'code.lineno' => source[1]
+            ) do
+              super(commands, redis_config, &block)
+            end
+          end
+        end
+
+        mod
+      end
+      private_class_method :build_patch_module
+    end
+  end
+end

data/lib/rails_otel_context/adapters/trilogy.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module RailsOtelContext
+  module Adapters
+    module Trilogy
+      module_function
+
+      def install!(app_root:, threshold_ms:)
+        return unless defined?(::Trilogy)
+
+        patch_module = patch_module_for
+        patch_module.configure(app_root: app_root, threshold_ms: threshold_ms)
+
+        return if ::Trilogy.ancestors.include?(patch_module)
+
+        ::Trilogy.prepend(patch_module)
+      end
+
+      def patch_module_for
+        @patch_module ||= build_patch_module
+      end
+
+      def build_patch_module
+        mod = Module.new do
+          class << self
+            include RailsOtelContext::SourceLocation
+
+            attr_accessor :app_root, :threshold_ms
+
+            def configure(app_root:, threshold_ms:)
+              @app_root = app_root.to_s
+              @threshold_ms = threshold_ms.to_f
+            end
+          end
+
+          # AR context attributes and span renaming are handled by
+          # CallContextProcessor.apply_db_context (via sql.active_record notification).
+          # This adapter only handles slow query source location tracking.
+          define_method(:query) do |sql|
+            source = mod.source_location_for_app
+            started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            result = super(sql)
+            duration_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at) * 1000.0
+
+            span = OpenTelemetry::Trace.current_span
+            if span.context.valid?
+              threshold = RailsOtelContext.configuration.trilogy_slow_query_threshold_ms
+              if source && duration_ms >= threshold
+                span.set_attribute('code.filepath', source[0])
+                span.set_attribute('code.lineno', source[1])
+                span.set_attribute('db.query.duration_ms', duration_ms.round(1))
+                span.set_attribute('db.query.slow_threshold_ms', threshold)
+              end
+            end
+
+            result
+          end
+        end
+
+        mod
+      end
+      private_class_method :build_patch_module
+    end
+  end
+end

data/lib/rails_otel_context/adapters.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require 'rails_otel_context/adapters/pg'
+require 'rails_otel_context/adapters/mysql2'
+require 'rails_otel_context/adapters/trilogy'
+require 'rails_otel_context/adapters/redis'
+require 'rails_otel_context/adapters/clickhouse'
+
+module RailsOtelContext
+  module Adapters
+    module_function
+
+    def install!(app_root:, config: RailsOtelContext.configuration)
+      install_pg!(app_root: app_root, config: config)
+      install_mysql2!(app_root: app_root, config: config)
+      install_trilogy!(app_root: app_root, config: config)
+      install_redis!(app_root: app_root, config: config)
+      install_clickhouse!(app_root: app_root, config: config)
+    end
+
+    def install_pg!(app_root:, config:)
+      return unless config.pg_slow_query_enabled
+
+      PG.install!(app_root: app_root, threshold_ms: config.pg_slow_query_threshold_ms)
+    end
+
+    def install_mysql2!(app_root:, config:)
+      return unless config.mysql2_slow_query_enabled
+
+      Mysql2.install!(app_root: app_root, threshold_ms: config.mysql2_slow_query_threshold_ms)
+    end
+
+    def install_trilogy!(app_root:, config:)
+      return unless config.trilogy_slow_query_enabled
+
+      Trilogy.install!(app_root: app_root, threshold_ms: config.trilogy_slow_query_threshold_ms)
+    end
+
+    def install_redis!(app_root:, config:)
+      return unless config.redis_source_enabled
+
+      Redis.install!(app_root: app_root)
+    end
+
+    def install_clickhouse!(app_root:, config:)
+      return unless config.clickhouse_enabled
+
+      Clickhouse.install!(app_root: app_root, threshold_ms: config.clickhouse_slow_query_threshold_ms)
+    end
+  end
+end