Add latest changes from gitlab-org/gitlab@master

1 files changed, 42 insertions, 7 deletions

diff --git a/spec/support/helpers/query_recorder.rb b/spec/support/helpers/query_recorder.rb
index 1d04014c9a6..fd200a1abf3 100644
--- a/spec/support/helpers/query_recorder.rb
+++ b/spec/support/helpers/query_recorder.rb
@@ -2,12 +2,15 @@
 
 module ActiveRecord
   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 = []
       @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
       ActiveRecord::Base.connection.unprepared_statement do
         ActiveSupport::Notifications.subscribed(method(:callback), 'sql.active_record', &block)
@@ -19,30 +22,62 @@ module ActiveRecord
       Gitlab::BacktraceCleaner.clean_backtrace(caller).each { |line| Rails.logger.debug("   --> #{line}") }
     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)
-      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
         @cached << values[:sql]
       elsif !values[:name]&.include?("SCHEMA")
         @log << values[:sql]
+        store_sql_by_source(values: values, backtrace: backtrace)
       end
     end
 
     def count
-      @log.count
+      @count ||= @log.count
     end
 
     def cached_count
-      @cached.count
+      @cached_count ||= @cached.count
     end
 
     def log_message
-      @log.join("\n\n")
+      @log_message ||= @log.join("\n\n")
     end
 
     def occurrences
-      @log.group_by(&:to_s).transform_values(&:count)
+      @occurrences ||= @log.group_by(&:to_s).transform_values(&:count)
     end
   end
 end