fin blog embed

active_record_query_counter.rb

# frozen_string_literal: true

class ActiveRecordQueryCounter
  # We want only one of these active
  include Singleton

  attr_reader :query_count

  def initialize
    reset
    set_trace_lines(1)
    _subscribe_for_count
    @cleaner = _backtrace_cleaner
  end

  private def _backtrace_cleaner
    # Remove unnecessary lines from the stack trace
    skip_re = %r{(initializers)|(lib/active_record_query_counter)|(lib/core_ext/active_record)}
    _cleaner = Rails.backtrace_cleaner.clone
    _cleaner.add_silencer { |line| skip_re.match(line) }
    _cleaner.add_filter { |line| line.sub("#{Rails.root}/", '') }
    _cleaner
  end

  public def reset
    @query_count = 0
  end

  public def set_trace_lines(num)
    @num_trace_lines = num
  end

  public def trace(num=nil)
    # provide a cleaned stack trace
    @cleaner.clean(caller).first(num || @num_trace_lines).map { |line| "[#{@query_count}] ↳ #{line}" }
  end

  private def _subscribe_for_count
    # Subscribe to query events so we can count them
    ActiveSupport::Notifications.unsubscribe(@count_subscriber) if @count_subscriber
    @count_subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |_name, _start, _finish, _id, payload|
      if payload[:name] != 'SCHEMA'
        @query_count += 1
      end
    end
  end
end