Commit b84a42ad authored by Ohad Dahan's avatar Ohad Dahan Committed by Robert Speicher

Enhance QueryRecorder to support finding queries by source + memoizing of read only attributes

parent 4c94998b
...@@ -33,6 +33,7 @@ module Quality ...@@ -33,6 +33,7 @@ module Quality
serializers serializers
services services
sidekiq sidekiq
support_specs
tasks tasks
uploaders uploaders
validators validators
......
...@@ -21,7 +21,7 @@ RSpec.describe Quality::TestLevel do ...@@ -21,7 +21,7 @@ RSpec.describe Quality::TestLevel do
context 'when level is unit' do context 'when level is unit' do
it 'returns a pattern' do it 'returns a pattern' do
expect(subject.pattern(:unit)) expect(subject.pattern(:unit))
.to eq("spec/{bin,config,db,dependencies,factories,finders,frontend,graphql,haml_lint,helpers,initializers,javascripts,lib,models,policies,presenters,rack_servers,replicators,routing,rubocop,serializers,services,sidekiq,tasks,uploaders,validators,views,workers,elastic_integration}{,/**/}*_spec.rb") .to eq("spec/{bin,config,db,dependencies,factories,finders,frontend,graphql,haml_lint,helpers,initializers,javascripts,lib,models,policies,presenters,rack_servers,replicators,routing,rubocop,serializers,services,sidekiq,support_specs,tasks,uploaders,validators,views,workers,elastic_integration}{,/**/}*_spec.rb")
end end
end end
...@@ -82,7 +82,7 @@ RSpec.describe Quality::TestLevel do ...@@ -82,7 +82,7 @@ RSpec.describe Quality::TestLevel do
context 'when level is unit' do context 'when level is unit' do
it 'returns a regexp' do it 'returns a regexp' do
expect(subject.regexp(:unit)) expect(subject.regexp(:unit))
.to eq(%r{spec/(bin|config|db|dependencies|factories|finders|frontend|graphql|haml_lint|helpers|initializers|javascripts|lib|models|policies|presenters|rack_servers|replicators|routing|rubocop|serializers|services|sidekiq|tasks|uploaders|validators|views|workers|elastic_integration)}) .to eq(%r{spec/(bin|config|db|dependencies|factories|finders|frontend|graphql|haml_lint|helpers|initializers|javascripts|lib|models|policies|presenters|rack_servers|replicators|routing|rubocop|serializers|services|sidekiq|support_specs|tasks|uploaders|validators|views|workers|elastic_integration)})
end end
end end
......
...@@ -2,12 +2,15 @@ ...@@ -2,12 +2,15 @@
module ActiveRecord module ActiveRecord
class QueryRecorder class QueryRecorder
attr_reader :log, :skip_cached, :cached attr_reader :log, :skip_cached, :cached, :data
UNKNOWN = %w(unknown unknown).freeze
def initialize(skip_cached: true, &block) def initialize(skip_cached: true, query_recorder_debug: false, &block)
@data = Hash.new { |h, k| h[k] = { count: 0, occurrences: [], backtrace: [] } }
@log = [] @log = []
@cached = [] @cached = []
@skip_cached = skip_cached @skip_cached = skip_cached
@query_recorder_debug = query_recorder_debug
# force replacement of bind parameters to give tests the ability to check for ids # force replacement of bind parameters to give tests the ability to check for ids
ActiveRecord::Base.connection.unprepared_statement do ActiveRecord::Base.connection.unprepared_statement do
ActiveSupport::Notifications.subscribed(method(:callback), 'sql.active_record', &block) ActiveSupport::Notifications.subscribed(method(:callback), 'sql.active_record', &block)
...@@ -19,30 +22,62 @@ module ActiveRecord ...@@ -19,30 +22,62 @@ module ActiveRecord
Gitlab::BacktraceCleaner.clean_backtrace(caller).each { |line| Rails.logger.debug(" --> #{line}") } Gitlab::BacktraceCleaner.clean_backtrace(caller).each { |line| Rails.logger.debug(" --> #{line}") }
end end
def get_sql_source(sql)
matches = sql.match(/,line:(?<line>.*):in\s+`(?<method>.*)'\*\//)
matches ? [matches[:line], matches[:method]] : UNKNOWN
end
def store_sql_by_source(values: {}, backtrace: nil)
full_name = get_sql_source(values[:sql]).join(':')
@data[full_name][:count] += 1
@data[full_name][:occurrences] << values[:sql]
@data[full_name][:backtrace] << backtrace
end
def find_query(query_regexp, limit, first_only: false)
out = []
@data.each_pair do |k, v|
if v[:count] > limit && k.match(query_regexp)
out << [k, v[:count]]
break if first_only
end
end
out.flatten! if first_only
out
end
def occurrences_by_line_method
@occurrences_by_line_method ||= @data.sort_by { |_, v| v[:count] }
end
def callback(name, start, finish, message_id, values) def callback(name, start, finish, message_id, values)
show_backtrace(values) if ENV['QUERY_RECORDER_DEBUG'] store_backtrace = ENV['QUERY_RECORDER_DEBUG'] || @query_recorder_debug
backtrace = store_backtrace ? show_backtrace(values) : nil
if values[:cached] && skip_cached if values[:cached] && skip_cached
@cached << values[:sql] @cached << values[:sql]
elsif !values[:name]&.include?("SCHEMA") elsif !values[:name]&.include?("SCHEMA")
@log << values[:sql] @log << values[:sql]
store_sql_by_source(values: values, backtrace: backtrace)
end end
end end
def count def count
@log.count @count ||= @log.count
end end
def cached_count def cached_count
@cached.count @cached_count ||= @cached.count
end end
def log_message def log_message
@log.join("\n\n") @log_message ||= @log.join("\n\n")
end end
def occurrences def occurrences
@log.group_by(&:to_s).transform_values(&:count) @occurrences ||= @log.group_by(&:to_s).transform_values(&:count)
end end
end end
end end
# frozen_string_literal: true
require 'spec_helper'
describe ActiveRecord::QueryRecorder do
class TestQueries < ActiveRecord::Base
self.table_name = 'schema_migrations'
end
describe 'detecting the right number of calls and their origin' do
it 'detects two separate queries' do
control = ActiveRecord::QueryRecorder.new query_recorder_debug: true do
2.times { TestQueries.count }
TestQueries.first
end
# Test first_only flag works as expected
expect(control.find_query(/.*query_recorder_spec.rb.*/, 0, first_only: true))
.to eq(control.find_query(/.*query_recorder_spec.rb.*/, 0).first)
# Check #find_query
expect(control.find_query(/.*/, 0).size)
.to eq(control.data.keys.size)
# Ensure exactly 2 COUNT queries were detected
expect(control.occurrences_by_line_method.last[1][:occurrences]
.find_all {|i| i.match(/SELECT COUNT/) }.count).to eq(2)
# Ensure exactly 1 LIMIT 1 (#first)
expect(control.occurrences_by_line_method.first[1][:occurrences]
.find_all { |i| i.match(/ORDER BY.*#{TestQueries.table_name}.*LIMIT 1/) }.count).to eq(1)
# Ensure 3 DB calls overall were executed
expect(control.log.size).to eq(3)
# Ensure memoization value match the raw value above
expect(control.count).to eq(control.log.size)
# Ensure we have only two sources of queries
expect(control.data.keys.size).to eq(2)
# Ensure we detect only queries from this file
expect(control.data.keys.find_all { |i| i.match(/query_recorder_spec.rb/) }.count).to eq(2)
end
end
end
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment