BigW Consortium Gitlab

query_recorder.rb 1.01 KB
Newer Older
1 2
module ActiveRecord
  class QueryRecorder
3
    attr_reader :log, :cached
4 5 6

    def initialize(&block)
      @log = []
7
      @cached = []
8 9 10 11
      ActiveSupport::Notifications.subscribed(method(:callback), 'sql.active_record', &block)
    end

    def callback(name, start, finish, message_id, values)
12 13 14 15 16
      if values[:name]&.include?("CACHE")
        @cached << values[:sql]
      elsif !values[:name]&.include?("SCHEMA")
        @log << values[:sql]
      end
17 18 19 20 21 22
    end

    def count
      @log.count
    end

23 24 25 26
    def cached_count
      @cached.count
    end

27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48
    def log_message
      @log.join("\n\n")
    end
  end
end

RSpec::Matchers.define :exceed_query_limit do |expected|
  supports_block_expectations

  match do |block|
    query_count(&block) > expected
  end

  failure_message_when_negated do |actual|
    "Expected a maximum of #{expected} queries, got #{@recorder.count}:\n\n#{@recorder.log_message}"
  end

  def query_count(&block)
    @recorder = ActiveRecord::QueryRecorder.new(&block)
    @recorder.count
  end
end