gitlab-org--gitlab-foss/spec/support/matchers/exceed_query_limit.rb

332 lines
7.3 KiB
Ruby

# frozen_string_literal: true
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
end
def for_query(query)
@query = query
self
end
def for_model(model)
table = model.table_name if model < ActiveRecord::Base
for_query(/(FROM|UPDATE|INSERT INTO|DELETE FROM)\s+"#{table}"/)
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)
query_recorder_count(expected)
else
expected
end
end
def actual_count
@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
# 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)
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
end
def skip_cached
true
end
def verify_count(&block)
@subject_block = block
actual_count > maximum
end
def maximum
expected_count + threshold
end
def failure_message
threshold_message = threshold > 0 ? " (+#{threshold})" : ''
counts = "#{expected_count}#{threshold_message}"
"Expected a maximum of #{counts} queries, got #{actual_count}:\n\n#{log_message}"
end
def strip_marginalia_annotations(logs)
logs.map { |log| log.sub(MARGINALIA_ANNOTATION_REGEX, '') }
end
end
RSpec::Matchers.define :issue_fewer_queries_than do
supports_block_expectations
include ExceedQueryLimitHelpers
def control
block_arg
end
def control_recorder
@control_recorder ||= ActiveRecord::QueryRecorder.new(&control)
end
def expected_count
control_recorder.count
end
def verify_count(&block)
@subject_block = block
# These blocks need to be evaluated in an expected order, in case
# the events in expected affect the counts in actual
expected_count
actual_count
actual_count < expected_count
end
match do |block|
verify_count(&block)
end
def failure_message
<<~MSG
Expected to issue fewer than #{expected_count} queries, but got #{actual_count}
#{log_message}
MSG
end
failure_message_when_negated do |actual|
<<~MSG
Expected query count of #{actual_count} to be less than #{expected_count}
#{log_message}
MSG
end
end
RSpec::Matchers.define :issue_same_number_of_queries_as do
supports_block_expectations
include ExceedQueryLimitHelpers
def control
block_arg
end
chain :or_fewer do
@or_fewer = true
end
chain :ignoring_cached_queries do
@skip_cached = true
end
def control_recorder
@control_recorder ||= ActiveRecord::QueryRecorder.new(&control)
end
def expected_count
control_recorder.count
end
def verify_count(&block)
@subject_block = block
# These blocks need to be evaluated in an expected order, in case
# the events in expected affect the counts in actual
expected_count
actual_count
if @or_fewer
actual_count <= expected_count
else
(expected_count - actual_count).abs <= threshold
end
end
match do |block|
verify_count(&block)
end
def failure_message
<<~MSG
Expected #{expected_count_message} queries, but got #{actual_count}
#{log_message}
MSG
end
failure_message_when_negated do |actual|
<<~MSG
Expected #{actual_count} not to equal #{expected_count_message}
#{log_message}
MSG
end
def expected_count_message
or_fewer_msg = "or fewer" if @or_fewer
threshold_msg = "(+/- #{threshold})" unless threshold == 0
["#{expected_count}", or_fewer_msg, threshold_msg].compact.join(' ')
end
def skip_cached
@skip_cached || false
end
end
RSpec::Matchers.define :exceed_all_query_limit do |expected|
supports_block_expectations
include ExceedQueryLimitHelpers
match do |block|
verify_count(&block)
end
failure_message_when_negated do |actual|
failure_message
end
def skip_cached
false
end
end
# Excludes cached methods from the query count
RSpec::Matchers.define :exceed_query_limit do |expected|
supports_block_expectations
include ExceedQueryLimitHelpers
match do |block|
verify_count(&block)
end
failure_message_when_negated do |actual|
failure_message
end
end