Merge branch 'fix-specs-with-query-analyzers' into 'master'

Allow nesting Query Analyzers

See merge request gitlab-org/gitlab!85531

lib/gitlab/database/query_analyzer.rb

@@ -30,52 +30,25 @@ def hook!
         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 @@ def begin!(analyzers = all_analyzers)
 
           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)

spec/lib/gitlab/database/query_analyzer_spec.rb

@@ -4,7 +4,7 @@
 
 RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
   let(:analyzer) { double(:query_analyzer) }
-  let(:user_analyzer) { double(:query_analyzer) }
+  let(:user_analyzer) { double(:user_query_analyzer) }
   let(:disabled_analyzer) { double(:disabled_query_analyzer) }
 
   before do
@@ -49,14 +49,36 @@
         end
       end
 
-      it 'does not evaluate enabled? again do yield block' do
-        expect(analyzer).not_to receive(:enabled?)
+      it 'does initialize analyzer only once' do
+        expect(analyzer).to receive(:enabled?).once
+        expect(analyzer).to receive(:begin!).once
+        expect(analyzer).to receive(:end!).once
 
         expect { |b| described_class.instance.within(&b) }.to yield_control
       end
 
-      it 'raises exception when trying to re-define analyzers' do
-        expect { |b| described_class.instance.within([user_analyzer], &b) }.to raise_error /Query analyzers are already defined, cannot re-define them/
+      it 'does initialize user analyzer when enabled' do
+        expect(user_analyzer).to receive(:enabled?).and_return(true)
+        expect(user_analyzer).to receive(:begin!)
+        expect(user_analyzer).to receive(:end!)
+
+        expect { |b| described_class.instance.within([user_analyzer], &b) }.to yield_control
+      end
+
+      it 'does initialize user analyzer only once' do
+        expect(user_analyzer).to receive(:enabled?).and_return(false, true)
+        expect(user_analyzer).to receive(:begin!).once
+        expect(user_analyzer).to receive(:end!).once
+
+        expect { |b| described_class.instance.within([user_analyzer, user_analyzer, user_analyzer], &b) }.to yield_control
+      end
+
+      it 'does not initializer user analyzer when disabled' do
+        expect(user_analyzer).to receive(:enabled?).and_return(false)
+        expect(user_analyzer).not_to receive(:begin!)
+        expect(user_analyzer).not_to receive(:end!)
+
+        expect { |b| described_class.instance.within([user_analyzer], &b) }.to yield_control
       end
     end
 
@@ -162,7 +184,7 @@
     def process_sql(sql)
       described_class.instance.within do
         ApplicationRecord.load_balancer.read_write do |connection|
-          described_class.instance.process_sql(sql, connection)
+          described_class.instance.send(:process_sql, sql, connection)
         end
       end
     end

spec/lib/gitlab/database/query_analyzers/gitlab_schemas_metrics_spec.rb

@@ -140,7 +140,7 @@
   def process_sql(model, sql)
     Gitlab::Database::QueryAnalyzer.instance.within do
       # Skip load balancer and retrieve connection assigned to model
-      Gitlab::Database::QueryAnalyzer.instance.process_sql(sql, model.retrieve_connection)
+      Gitlab::Database::QueryAnalyzer.instance.send(:process_sql, sql, model.retrieve_connection)
     end
   end
 end

spec/lib/gitlab/database/query_analyzers/restrict_allowed_schemas_spec.rb

@@ -155,7 +155,7 @@ def process_sql(sql, model = ActiveRecord::Base)
       yield if block_given?
 
       # Skip load balancer and retrieve connection assigned to model
-      Gitlab::Database::QueryAnalyzer.instance.process_sql(sql, model.retrieve_connection)
+      Gitlab::Database::QueryAnalyzer.instance.send(:process_sql, sql, model.retrieve_connection)
     end
   end
 end

spec/support/database/query_analyzer.rb

@@ -6,13 +6,17 @@
 RSpec.configure do |config|
   config.before do |example|
     if example.metadata.fetch(:query_analyzers, true)
-      ::Gitlab::Database::QueryAnalyzer.instance.begin!
+      ::Gitlab::Database::QueryAnalyzer.instance.begin!(
+        ::Gitlab::Database::QueryAnalyzer.instance.all_analyzers
+      )
     end
   end
 
   config.after do |example|
     if example.metadata.fetch(:query_analyzers, true)
-      ::Gitlab::Database::QueryAnalyzer.instance.end!
+      ::Gitlab::Database::QueryAnalyzer.instance.end!(
+        ::Gitlab::Database::QueryAnalyzer.instance.all_analyzers
+      )
     end
   end
 end