Allow exceed_query_limit to ignore irrelevant queries

This also improves the diff output so you can see where the extra
queries are actually coming from.

spec/requests/api/graphql/project_query_spec.rb

@@ -151,8 +151,9 @@ RSpec.describe 'getting project information' do
           )))
     end
 
-    it 'can lookahead to eliminate N+1 queries', :use_clean_rails_memory_store_caching, :request_store do
-      expect { run_query(10) }.to issue_same_number_of_queries_as { run_query(1) }.or_fewer.ignoring_cached_queries
+    it 'can lookahead to eliminate N+1 queries' do
+      baseline = ActiveRecord::QueryRecorder.new { run_query(1) }
+      expect { run_query(10) }.not_to exceed_query_limit(baseline)
     end
   end
 

spec/support/matchers/exceed_query_limit.rb

@@ -3,6 +3,13 @@
 module ExceedQueryLimitHelpers
   MARGINALIA_ANNOTATION_REGEX = %r{\s*\/\*.*\*\/}.freeze
 
+  DB_QUERY_RE = Regexp.union([
+    /^(?<prefix>SELECT .* FROM "?[a-z_]+"?) (?<suffix>.*)$/m,
+    /^(?<prefix>UPDATE "?[a-z_]+"?) (?<suffix>.*)$/m,
+    /^(?<prefix>INSERT INTO "[a-z_]+" \((?:"[a-z_]+",?\s?)+\)) (?<suffix>.*)$/m,
+    /^(?<prefix>DELETE FROM "[a-z_]+") (?<suffix>.*)$/m
+  ]).freeze
+
   def with_threshold(threshold)
     @threshold = threshold
     self
@@ -13,41 +20,129 @@ module ExceedQueryLimitHelpers
     self
   end
 
+  def show_common_queries
+    @show_common_queries = true
+    self
+  end
+
+  def ignoring(pattern)
+    @ignoring_pattern = pattern
+    self
+  end
+
   def threshold
     @threshold.to_i
   end
 
   def expected_count
     if expected.is_a?(ActiveRecord::QueryRecorder)
-      expected.count
+      query_recorder_count(expected)
     else
       expected
     end
   end
 
   def actual_count
-    @actual_count ||= if @query
-                        recorder.log.select { |recorded| recorded =~ @query }.size
-                      else
-                        recorder.count
-                      end
+    @actual_count ||= query_recorder_count(recorder)
+  end
+
+  def query_recorder_count(query_recorder)
+    return query_recorder.count unless @query || @ignoring_pattern
+
+    query_log(query_recorder).size
+  end
+
+  def query_log(query_recorder)
+    filtered = query_recorder.log
+    filtered = filtered.select { |q| q =~ @query } if @query
+    filtered = filtered.reject { |q| q =~ @ignoring_pattern } if @ignoring_pattern
+    filtered
   end
 
   def recorder
     @recorder ||= ActiveRecord::QueryRecorder.new(skip_cached: skip_cached, &@subject_block)
   end
 
-  def count_queries(queries)
-    queries.each_with_object(Hash.new(0)) { |query, counts| counts[query] += 1 }
+  # Take a query recorder and tabulate the frequencies of suffixes for each prefix.
+  #
+  # @return Hash[String, Hash[String, Int]]
+  #
+  # Example:
+  #
+  # r = ActiveRecord::QueryRecorder.new do
+  #   SomeTable.create(x: 1, y: 2, z: 3)
+  #   SomeOtherTable.where(id: 1).first
+  #   SomeTable.create(x: 4, y: 5, z: 6)
+  #   SomeOtherTable.all
+  # end
+  # count_queries(r)
+  # #=>
+  #  {
+  #    'INSERT INTO "some_table" VALUES' => {
+  #      '(1,2,3)' => 1,
+  #      '(4,5,6)' => 1
+  #    },
+  #    'SELECT * FROM "some_other_table"' => {
+  #      'WHERE id = 1 LIMIT 1' => 1,
+  #      '' => 2
+  #    }
+  #  }
+  def count_queries(query_recorder)
+    strip_marginalia_annotations(query_log(query_recorder))
+      .map { |q| query_group_key(q) }
+      .group_by { |k| k[:prefix] }
+      .transform_values { |keys| frequencies(:suffix, keys) }
+  end
+
+  def frequencies(key, things)
+    things.group_by { |x| x[key] }.transform_values(&:size)
+  end
+
+  def query_group_key(query)
+    DB_QUERY_RE.match(query) || { prefix: query, suffix: '' }
+  end
+
+  def diff_query_counts(expected, actual)
+    expected_counts = expected.transform_values do |suffixes|
+      suffixes.transform_values { |n| [n, 0] }
+    end
+    recorded_counts = actual.transform_values do |suffixes|
+      suffixes.transform_values { |n| [0, n] }
+    end
+
+    combined_counts = expected_counts.merge(recorded_counts) do |_k, exp, got|
+      exp.merge(got) do |_k, exp_counts, got_counts|
+        exp_counts.zip(got_counts).map { |a, b| a + b }
+      end
+    end
+
+    unless @show_common_queries
+      combined_counts = combined_counts.transform_values do |suffs|
+        suffs.reject { |_k, counts| counts.first == counts.second }
+      end
+    end
+
+    combined_counts.reject { |_prefix, suffs| suffs.empty? }
+  end
+
+  def diff_query_group_message(query, suffixes)
+    suffix_messages = suffixes.map do |s, counts|
+      "-- (expected: #{counts.first}, got: #{counts.second})\n   #{s}"
+    end
+
+    "#{query}...\n#{suffix_messages.join("\n")}"
   end
 
   def log_message
     if expected.is_a?(ActiveRecord::QueryRecorder)
-      counts = count_queries(strip_marginalia_annotations(expected.log))
-      extra_queries = strip_marginalia_annotations(@recorder.log).reject { |query| counts[query] -= 1 unless counts[query] == 0 }
-      extra_queries_display = count_queries(extra_queries).map { |query, count| "[#{count}] #{query}" }
-
-      (['Extra queries:'] + extra_queries_display).join("\n\n")
+      diff_counts = diff_query_counts(count_queries(expected), count_queries(@recorder))
+      sections = diff_counts.map { |q, suffixes| diff_query_group_message(q, suffixes) }
+
+      <<~MSG
+      Query Diff:
+      -----------
+      #{sections.join("\n\n")}
+      MSG
     else
       @recorder.log_message
     end

spec/support_specs/matchers/exceed_query_limit_helpers_spec.rb

@@ -22,6 +22,232 @@ RSpec.describe ExceedQueryLimitHelpers do
     end
   end
 
+  describe '#diff_query_group_message' do
+    it 'prints a group helpfully' do
+      test_matcher = TestMatcher.new
+      suffixes = {
+        'WHERE x = z' => [1, 1],
+        'WHERE x = y' => [1, 2],
+        'LIMIT 1' => [1, 0]
+      }
+
+      message = test_matcher.diff_query_group_message('SELECT * FROM foo', suffixes)
+
+      expect(message).to eq(<<~MSG.chomp)
+      SELECT * FROM foo...
+      -- (expected: 1, got: 1)
+         WHERE x = z
+      -- (expected: 1, got: 2)
+         WHERE x = y
+      -- (expected: 1, got: 0)
+         LIMIT 1
+      MSG
+    end
+  end
+
+  describe '#diff_query_counts' do
+    let(:expected) do
+      ActiveRecord::QueryRecorder.new do
+        TestQueries.where(version: 'foobar').to_a
+        TestQueries.where(version: 'also foobar and baz').to_a
+        TestQueries.count
+        TestQueries.first
+        TestQueries.where(version: 'foobar').to_a
+        TestQueries.where(version: 'x').update_all(version: 'y')
+        TestQueries.where(version: 'foobar').count
+        TestQueries.where(version: 'z').delete_all
+      end
+    end
+
+    let(:actual) do
+      ActiveRecord::QueryRecorder.new do
+        TestQueries.where(version: 'foobar').to_a
+        TestQueries.where(version: 'also foobar and baz').to_a
+        TestQueries.count
+        TestQueries.create!(version: 'x')
+        TestQueries.where(version: 'foobar').to_a
+        TestQueries.where(version: 'x').update_all(version: 'y')
+        TestQueries.where(version: 'foobar').count
+        TestQueries.count
+        TestQueries.where(version: 'y').update_all(version: 'z')
+        TestQueries.where(version: 'z').delete_all
+      end
+    end
+
+    it 'merges two query counts' do
+      test_matcher = TestMatcher.new
+
+      diff = test_matcher.diff_query_counts(
+        test_matcher.count_queries(expected),
+        test_matcher.count_queries(actual)
+      )
+
+      expect(diff).to eq({
+        "SELECT \"schema_migrations\".* FROM \"schema_migrations\"" => {
+          "ORDER BY \"schema_migrations\".\"version\" ASC LIMIT 1" => [1, 0]
+        },
+        "SELECT COUNT(*) FROM \"schema_migrations\"" => { "" => [1, 2] },
+        "UPDATE \"schema_migrations\"" => {
+          "SET \"version\" = 'z' WHERE \"schema_migrations\".\"version\" = 'y'" => [0, 1]
+        },
+        "SAVEPOINT active_record_1" => { "" => [0, 1] },
+        "INSERT INTO \"schema_migrations\" (\"version\")" => {
+          "VALUES ('x') RETURNING \"version\"" => [0, 1]
+        },
+        "RELEASE SAVEPOINT active_record_1" => { "" => [0, 1] }
+      })
+    end
+
+    it 'can show common queries if so desired' do
+      test_matcher = TestMatcher.new.show_common_queries
+
+      diff = test_matcher.diff_query_counts(
+        test_matcher.count_queries(expected),
+        test_matcher.count_queries(actual)
+      )
+
+      expect(diff).to eq({
+        "SELECT \"schema_migrations\".* FROM \"schema_migrations\"" => {
+          "WHERE \"schema_migrations\".\"version\" = 'foobar'" => [2, 2],
+          "WHERE \"schema_migrations\".\"version\" = 'also foobar and baz'" => [1, 1],
+          "ORDER BY \"schema_migrations\".\"version\" ASC LIMIT 1" => [1, 0]
+        },
+        "SELECT COUNT(*) FROM \"schema_migrations\"" => {
+          "" => [1, 2],
+          "WHERE \"schema_migrations\".\"version\" = 'foobar'" => [1, 1]
+        },
+        "UPDATE \"schema_migrations\"" => {
+          "SET \"version\" = 'y' WHERE \"schema_migrations\".\"version\" = 'x'" => [1, 1],
+          "SET \"version\" = 'z' WHERE \"schema_migrations\".\"version\" = 'y'" => [0, 1]
+        },
+        "DELETE FROM \"schema_migrations\"" => {
+          "WHERE \"schema_migrations\".\"version\" = 'z'" => [1, 1]
+        },
+        "SAVEPOINT active_record_1" => {
+          "" => [0, 1]
+        },
+        "INSERT INTO \"schema_migrations\" (\"version\")" => {
+          "VALUES ('x') RETURNING \"version\"" => [0, 1]
+        },
+        "RELEASE SAVEPOINT active_record_1" => {
+          "" => [0, 1]
+        }
+      })
+    end
+  end
+
+  describe '#count_queries' do
+    it 'handles queries with suffixes over multiple lines' do
+      test_matcher = TestMatcher.new
+
+      recorder = ActiveRecord::QueryRecorder.new do
+        TestQueries.find_by(version: %w(foo bar baz).join("\n"))
+        TestQueries.find_by(version: %w(foo biz baz).join("\n"))
+        TestQueries.find_by(version: %w(foo bar baz).join("\n"))
+      end
+
+      recorder.count
+
+      expect(test_matcher.count_queries(recorder)).to eq({
+        'SELECT "schema_migrations".* FROM "schema_migrations"' => {
+          %Q[WHERE "schema_migrations"."version" = 'foo\nbar\nbaz' LIMIT 1] => 2,
+          %Q[WHERE "schema_migrations"."version" = 'foo\nbiz\nbaz' LIMIT 1] => 1
+        }
+      })
+    end
+
+    it 'can aggregate queries' do
+      test_matcher = TestMatcher.new
+
+      recorder = ActiveRecord::QueryRecorder.new do
+        TestQueries.where(version: 'foobar').to_a
+        TestQueries.where(version: 'also foobar and baz').to_a
+        TestQueries.count
+        TestQueries.create!(version: 'x')
+        TestQueries.first
+        TestQueries.where(version: 'foobar').to_a
+        TestQueries.where(version: 'x').update_all(version: 'y')
+        TestQueries.where(version: 'foobar').count
+        TestQueries.count
+        TestQueries.where(version: 'y').update_all(version: 'z')
+        TestQueries.where(version: 'z').delete_all
+      end
+
+      recorder.count
+
+      expect(test_matcher.count_queries(recorder)).to eq({
+        'SELECT "schema_migrations".* FROM "schema_migrations"' => {
+          %q[WHERE "schema_migrations"."version" = 'foobar'] => 2,
+          %q[WHERE "schema_migrations"."version" = 'also foobar and baz'] => 1,
+          %q[ORDER BY "schema_migrations"."version" ASC LIMIT 1] => 1
+        },
+        'SELECT COUNT(*) FROM "schema_migrations"' => {
+          "" => 2,
+          %q[WHERE "schema_migrations"."version" = 'foobar'] => 1
+        },
+        'SAVEPOINT active_record_1' => { "" => 1 },
+        'INSERT INTO "schema_migrations" ("version")' => {
+          %q[VALUES ('x') RETURNING "version"] => 1
+        },
+        'RELEASE SAVEPOINT active_record_1' => { "" => 1 },
+        'UPDATE "schema_migrations"' => {
+          %q[SET "version" = 'y' WHERE "schema_migrations"."version" = 'x'] => 1,
+          %q[SET "version" = 'z' WHERE "schema_migrations"."version" = 'y'] => 1
+        },
+        'DELETE FROM "schema_migrations"' => {
+          %q[WHERE "schema_migrations"."version" = 'z'] => 1
+        }
+      })
+    end
+  end
+
+  it 'can count queries' do
+    test_matcher = TestMatcher.new
+    test_matcher.verify_count do
+      TestQueries.where(version: 'foobar').to_a
+      TestQueries.where(version: 'also foobar and baz').to_a
+      TestQueries.first
+      TestQueries.count
+    end
+
+    expect(test_matcher.actual_count).to eq(4)
+  end
+
+  it 'can select specific queries' do
+    test_matcher = TestMatcher.new.for_query(/foobar/)
+    test_matcher.verify_count do
+      TestQueries.where(version: 'foobar').to_a
+      TestQueries.where(version: 'also foobar and baz').to_a
+      TestQueries.first
+      TestQueries.count
+    end
+
+    expect(test_matcher.actual_count).to eq(2)
+  end
+
+  it 'can ignore specific queries' do
+    test_matcher = TestMatcher.new.ignoring(/foobar/)
+    test_matcher.verify_count do
+      TestQueries.where(version: 'foobar').to_a
+      TestQueries.where(version: 'also foobar and baz').to_a
+      TestQueries.first
+    end
+
+    expect(test_matcher.actual_count).to eq(1)
+  end
+
+  it 'can perform inclusion and exclusion' do
+    test_matcher = TestMatcher.new.for_query(/foobar/).ignoring(/baz/)
+    test_matcher.verify_count do
+      TestQueries.where(version: 'foobar').to_a
+      TestQueries.where(version: 'also foobar and baz').to_a
+      TestQueries.first
+      TestQueries.count
+    end
+
+    expect(test_matcher.actual_count).to eq(1)
+  end
+
   it 'does not contain marginalia annotations' do
     test_matcher = TestMatcher.new
     test_matcher.verify_count do