Counting the number of queries performed

Question

I\'d like to test that a certain piece of code performs as few SQL queries as possible.

ActiveRecord::TestCase seems to have its own assert_querie

Answer 1

I added the ability to check queries per table based on Yuriy's solution

# spec/support/query_counter.rb
require 'support/matchers/query_limit'

module ActiveRecord
  class QueryCounter
    attr_reader :queries

    def initialize
      @queries = Hash.new 0
    end

    def to_proc
      lambda(&method(:callback))
    end

    def callback(name, start, finish, message_id, values)
      sql = values[:sql]

      if sql.include? 'SAVEPOINT'
        table = :savepoints
      else
        finder = /select.+"(.+)"\..+from/i if sql.include? 'SELECT'
        finder = /insert.+"(.+)".\(/i if sql.include? 'INSERT'
        finder = /update.+"(.+)".+set/i if sql.include? 'UPDATE'
        finder = /delete.+"(.+)" where/i if sql.include? 'DELETE'
        table = sql.match(finder)&.send(:[],1)&.to_sym
      end

      @queries[table] += 1 unless %w(CACHE SCHEMA).include?(values[:name])

      return @queries
    end

    def query_count(table = nil)
      if table
        @queries[table]
      else
        @queries.values.sum
      end
    end
  end
end

The RSpec matchers look like

# spec/support/matchers/query_limit.rb
RSpec::Matchers.define :exceed_query_limit do |expected, table|
  supports_block_expectations

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

  def query_count(table, &block)
    @counter = ActiveRecord::QueryCounter.new
    ActiveSupport::Notifications.subscribed(@counter.to_proc, 'sql.active_record', &block)
    @counter.query_count table
  end

  failure_message_when_negated do |actual|
    queries = 'query'.pluralize expected
    table_name = table.to_s.singularize.humanize.downcase if table

    out = "expected to run a maximum of #{expected}"
    out += " #{table_name}" if table
    out += " #{queries}, but got #{@counter.query_count table}"
  end
end

RSpec::Matchers.define :meet_query_limit do |expected, table|
  supports_block_expectations

  match do |block|
    if expected.is_a? Hash
      results = queries_count(table, &block)
      expected.all? { |table, count| results[table] == count }
    else
      query_count(&block) == expected
    end
  end

  def queries_count(table, &block)
    @counter = ActiveRecord::QueryCounter.new
    ActiveSupport::Notifications.subscribed(@counter.to_proc, 'sql.active_record', &block)
    @counter.queries
  end

  def query_count(&block)
    @counter = ActiveRecord::QueryCounter.new
    ActiveSupport::Notifications.subscribed(@counter.to_proc, 'sql.active_record', &block)
    @counter.query_count
  end

  def message(expected, table, negated = false)
    queries = 'query'.pluralize expected
    if expected.is_a? Hash
      results = @counter.queries
      table, expected = expected.find { |table, count| results[table] != count }
    end

    table_name = table.to_s.singularize.humanize.downcase if table

    out = 'expected to'
    out += ' not' if negated
    out += " run exactly #{expected}"
    out += " #{table_name}" if table
    out += " #{queries}, but got #{@counter.query_count table}"
  end

  failure_message do |actual|
    message expected, table
  end

  failure_message_when_negated do |actual|
    message expected, table, true
  end
end

Usage

expect { MyModel.do_the_queries }.to_not meet_query_limit(3)
expect { MyModel.do_the_queries }.to meet_query_limit(3)
expect { MyModel.do_the_queries }.to meet_query_limit(my_models: 2, other_tables: 1)

Answer 2

I think you answered your own question by mentioning assert_queries, but here goes:

I would recommend taking a look at the code behind assert_queries and using that to build your own method which you can use to count queries. The main magic involved here is this line:

ActiveSupport::Notifications.subscribe('sql.active_record', SQLCounter.new)

I had a bit of a tinker this morning and ripped out the parts of ActiveRecord that do the query counting and came up with this:

module ActiveRecord
  class QueryCounter
    cattr_accessor :query_count do
      0
    end

    IGNORED_SQL = [/^PRAGMA (?!(table_info))/, /^SELECT currval/, /^SELECT CAST/, /^SELECT @@IDENTITY/, /^SELECT @@ROWCOUNT/, /^SAVEPOINT/, /^ROLLBACK TO SAVEPOINT/, /^RELEASE SAVEPOINT/, /^SHOW max_identifier_length/]

    def call(name, start, finish, message_id, values)
      # FIXME: this seems bad. we should probably have a better way to indicate
      # the query was cached
      unless 'CACHE' == values[:name]
        self.class.query_count += 1 unless IGNORED_SQL.any? { |r| values[:sql] =~ r }
      end
    end
  end
end

ActiveSupport::Notifications.subscribe('sql.active_record', ActiveRecord::QueryCounter.new)

module ActiveRecord
  class Base
    def self.count_queries(&block)
      ActiveRecord::QueryCounter.query_count = 0
      yield
      ActiveRecord::QueryCounter.query_count
    end
  end
end

You will be able to reference the ActiveRecord::Base.count_queries method anywhere. Pass it a block wherein your queries are run and it will return the number of queries that have been executed:

ActiveRecord::Base.count_queries do
  Ticket.first
end

Returns "1" for me. To make this work: put it in a file at lib/active_record/query_counter.rb and require it in your config/application.rb file like this:

require 'active_record/query_counter'

Hey presto!

---

A little bit of explanation probably is required. When we call this line:

    ActiveSupport::Notifications.subscribe('sql.active_record', ActiveRecord::QueryCounter.new)

We hook into Rails 3's little notifications framework. It's a shiny little addition to the latest major version of Rails that nobody really knows about. It allows us to subscribe to notifications of events within Rails by using the subscribe method. We pass in the event we want to subscribe to as the first argument then any object that responds to call as the second.

In this case when a query is executed our little query counter will dutifully increment the ActiveRecord::QueryCounter.query_count variable, but only for the real queries.

Anyway, this was fun. I hope it comes useful to you.

Answer 3

Here's a version that makes it easy to count queries matching a given pattern.

module QueryCounter

  def self.count_selects(&block)
    count(pattern: /^(\s+)?SELECT/, &block)
  end

  def self.count(pattern: /(.*?)/, &block)
    counter = 0

    callback = ->(name, started, finished, callback_id, payload) {
      counter += 1 if payload[:sql].match(pattern)
      # puts "match? #{!!payload[:sql].match(pattern)}: #{payload[:sql]}"
    }

    # http://api.rubyonrails.org/classes/ActiveSupport/Notifications.html
    ActiveSupport::Notifications.subscribed(callback, "sql.active_record", &block)

    counter
  end

end

Usage:

test "something" do
  query_count = count_selects {
    Thing.first
    Thing.create!(size: "huge")
  }
  assert_equal 1, query_count
end

Answer 4

I ended up creating a tiny gem to abstract this problem: sql_spy.

Just add it to your Gemfile:

gem "sql_spy"

Wrap your code inside SqlSpy.track { ... }:

queries = SqlSpy.track do
  # Some code that triggers ActiveRecord queries
  users = User.all
  posts = BlogPost.all
end

...and use the return value of the block in your assertions:

expect(queries.size).to eq(2)
expect(queries[0].sql).to eq("SELECT * FROM users;")
expect(queries[0].model_name).to eq("User")
expect(queries[0].select?).to be_true
expect(queries[0].duration).to eq(1.5)

Answer 5

• helpful error message
• removes subscribers after execution

(based on Jaime Cham's answer)

class ActiveSupport::TestCase
  def sql_queries(&block)
    queries = []
    counter = ->(*, payload) {
      queries << payload.fetch(:sql) unless ["CACHE", "SCHEMA"].include?(payload.fetch(:name))
    }

    ActiveSupport::Notifications.subscribed(counter, "sql.active_record", &block)

    queries
  end

  def assert_sql_queries(expected, &block)
    queries = sql_queries(&block)
    queries.count.must_equal(
      expected,
      "Expected #{expected} queries, but found #{queries.count}:\n#{queries.join("\n")}"
    )
  end
end

Answer 6

Here's another formulation of Ryan's and Yuriy's solution that's just a function you add to your test_helper.rb:

def count_queries &block
  count = 0

  counter_f = ->(name, started, finished, unique_id, payload) {
    unless payload[:name].in? %w[ CACHE SCHEMA ]
      count += 1
    end
  }

  ActiveSupport::Notifications.subscribed(counter_f, "sql.active_record", &block)

  count
end

Usage is just:

c = count_queries do
  SomeModel.first
end