1 files changed, 32 insertions, 44 deletions

diff --git a/lib/gitlab/database/query_analyzer.rb b/lib/gitlab/database/query_analyzer.rb
index 0c78dda734c..6f64d04270f 100644
--- a/lib/gitlab/database/query_analyzer.rb
+++ b/lib/gitlab/database/query_analyzer.rb
@@ -30,52 +30,25 @@ module Gitlab
         end
       end
 
-      def within(user_analyzers = nil)
-        # Due to singleton nature of analyzers
-        # only an outer invocation of the `.within`
-        # is allowed to initialize them
-        if already_within?
-          raise 'Query analyzers are already defined, cannot re-define them.' if user_analyzers
-
-          return yield
-        end
-
-        begin!(user_analyzers || all_analyzers)
+      def within(analyzers = all_analyzers)
+        newly_enabled_analyzers = begin!(analyzers)
 
         begin
           yield
         ensure
-          end!
+          end!(newly_enabled_analyzers)
         end
       end
 
-      def already_within?
-        # If analyzers are set they are already configured
-        !enabled_analyzers.nil?
-      end
+      # Enable query analyzers (only the ones that were not yet enabled)
+      # Returns a list of newly enabled analyzers
+      def begin!(analyzers)
+        analyzers.select do |analyzer|
+          next if enabled_analyzers.include?(analyzer)
 
-      def process_sql(sql, connection)
-        analyzers = enabled_analyzers
-        return unless analyzers&.any?
-
-        parsed = parse(sql, connection)
-        return unless parsed
-
-        analyzers.each do |analyzer|
-          next if analyzer.suppressed? && !analyzer.requires_tracking?(parsed)
-
-          analyzer.analyze(parsed)
-        rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
-          # We catch all standard errors to prevent validation errors to introduce fatal errors in production
-          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
-        end
-      end
-
-      # Enable query analyzers
-      def begin!(analyzers = all_analyzers)
-        analyzers = analyzers.select do |analyzer|
           if analyzer.enabled?
             analyzer.begin!
+            enabled_analyzers.append(analyzer)
 
             true
           end
@@ -84,25 +57,40 @@ module Gitlab
 
           false
         end
-
-        Thread.current[:query_analyzer_enabled_analyzers] = analyzers
       end
 
-      # Disable enabled query analyzers
-      def end!
-        enabled_analyzers.select do |analyzer|
+      # Disable enabled query analyzers (only the ones that were enabled previously)
+      def end!(analyzers)
+        analyzers.each do |analyzer|
+          next unless enabled_analyzers.delete(analyzer)
+
           analyzer.end!
         rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
           Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
         end
-
-        Thread.current[:query_analyzer_enabled_analyzers] = nil
       end
 
       private
 
       def enabled_analyzers
-        Thread.current[:query_analyzer_enabled_analyzers]
+        Thread.current[:query_analyzer_enabled_analyzers] ||= []
+      end
+
+      def process_sql(sql, connection)
+        analyzers = enabled_analyzers
+        return unless analyzers&.any?
+
+        parsed = parse(sql, connection)
+        return unless parsed
+
+        analyzers.each do |analyzer|
+          next if analyzer.suppressed? && !analyzer.requires_tracking?(parsed)
+
+          analyzer.analyze(parsed)
+        rescue StandardError, ::Gitlab::Database::QueryAnalyzers::Base::QueryAnalyzerError => e
+          # We catch all standard errors to prevent validation errors to introduce fatal errors in production
+          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
+        end
       end
 
       def parse(sql, connection)