Add `QueryAnalyzer::Base.suppressed?` method

This allows to temporarily suppress the analyzer
for any purpose in a consistent way.

lib/gitlab/database/query_analyzer.rb

@@ -58,6 +58,8 @@ module Gitlab
         return unless parsed
 
         analyzers.each do |analyzer|
+          next if analyzer.suppressed?
+
           analyzer.analyze(parsed)
         rescue StandardError => e
           # We catch all standard errors to prevent validation errors to introduce fatal errors in production

lib/gitlab/database/query_analyzers/base.rb

@@ -4,16 +4,32 @@ module Gitlab
   module Database
     module QueryAnalyzers
       class Base
+        def self.suppressed?
+          Thread.current[self.suppress_key]
+        end
+
+        def self.suppress=(value)
+          Thread.current[self.suppress_key] = value
+        end
+
+        def self.with_suppressed(value = true, &blk)
+          previous = self.suppressed?
+          self.suppress = value
+          yield
+        ensure
+          self.suppress = previous
+        end
+
         def self.begin!
-          Thread.current[self.class.name] = {}
+          Thread.current[self.context_key] = {}
         end
 
         def self.end!
-          Thread.current[self.class.name] = nil
+          Thread.current[self.context_key] = nil
         end
 
         def self.context
-          Thread.current[self.class.name]
+          Thread.current[self.context_key]
         end
 
         def self.enabled?
@@ -23,6 +39,14 @@ module Gitlab
         def self.analyze(parsed)
           raise NotImplementedError
         end
+
+        def self.context_key
+          "#{self.class.name}_context"
+        end
+
+        def self.suppress_key
+          "#{self.class.name}_suppressed"
+        end
       end
     end
   end

lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb

@@ -6,23 +6,6 @@ module Gitlab
       class PreventCrossDatabaseModification < Database::QueryAnalyzers::Base
         CrossDatabaseModificationAcrossUnsupportedTablesError = Class.new(StandardError)
 
-        def self.allow_cross_database_modification?
-          Thread.current[:prevent_cross_database_modification_allowed]
-        end
-
-        def self.allow_cross_database_modification=(value)
-          Thread.current[:prevent_cross_database_modification_allowed] = value
-        end
-
-        def self.with_allow_cross_database_modification(value, &blk)
-          previous = self.allow_cross_database_modification?
-          self.allow_cross_database_modification = value
-
-          yield
-        ensure
-          self.allow_cross_database_modification = previous
-        end
-
         # This method will allow cross database modifications within the block
         # Example:
         #
@@ -30,13 +13,13 @@ module Gitlab
         #   create(:build) # inserts ci_build and project record in one transaction
         # end
         def self.allow_cross_database_modification_within_transaction(url:, &blk)
-          self.with_allow_cross_database_modification(true, &blk)
+          self.with_suppressed(true, &blk)
         end
 
         # This method will prevent cross database modifications within the block
         # if it was allowed previously
         def self.with_cross_database_modification_prevented(&blk)
-          self.with_allow_cross_database_modification(false, &blk)
+          self.with_suppressed(false, &blk)
         end
 
         def self.begin!
@@ -55,7 +38,6 @@ module Gitlab
 
         # rubocop:disable Metrics/AbcSize
         def self.analyze(parsed)
-          return if self.allow_cross_database_modification?
           return if in_factory_bot_create?
 
           database = ::Gitlab::Database.db_config_name(parsed.connection)
@@ -112,10 +94,6 @@ module Gitlab
         rescue CrossDatabaseModificationAcrossUnsupportedTablesError => e
           ::Gitlab::ErrorTracking.track_exception(e, { gitlab_schemas: schemas, tables: all_tables, query: parsed.sql })
           raise if raise_exception?
-        rescue StandardError => e
-          # Extra safety net to ensure we never raise in production
-          # if something goes wrong in this logic
-          Gitlab::ErrorTracking.track_and_raise_for_dev_exception(e)
         end
         # rubocop:enable Metrics/AbcSize
 

spec/lib/gitlab/database/query_analyzer_spec.rb

@@ -9,6 +9,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
   before do
     allow(described_class.instance).to receive(:all_analyzers).and_return([analyzer, disabled_analyzer])
     allow(analyzer).to receive(:enabled?).and_return(true)
+    allow(analyzer).to receive(:suppressed?).and_return(false)
     allow(analyzer).to receive(:begin!)
     allow(analyzer).to receive(:end!)
     allow(disabled_analyzer).to receive(:enabled?).and_return(false)
@@ -125,6 +126,13 @@ RSpec.describe Gitlab::Database::QueryAnalyzer, query_analyzers: false do
       expect { process_sql("SELECT 1 FROM projects") }.not_to raise_error
     end
 
+    it 'does not call analyze on suppressed analyzers' do
+      expect(analyzer).to receive(:suppressed?).and_return(true)
+      expect(analyzer).not_to receive(:analyze)
+
+      expect { process_sql("SELECT 1 FROM projects") }.not_to raise_error
+    end
+
     def process_sql(sql)
       described_class.instance.within do
         ApplicationRecord.load_balancer.read_write do |connection|

spec/support/database/prevent_cross_database_modification.rb

@@ -11,14 +11,21 @@ CROSS_DB_MODIFICATION_ALLOW_LIST = Set.new(YAML.load_file(File.join(__dir__, 'cr
 RSpec.configure do |config|
   config.include(PreventCrossDatabaseModificationSpecHelpers)
 
+  # By default allow cross-modifications as we want to observe only transactions
+  # within a specific block of execution which is defined be `before(:each)` and `after(:each)`
+  config.before(:all) do
+    ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.suppress = true
+  end
+
   # Using before and after blocks because the around block causes problems with the let_it_be
   # record creations. It makes an extra savepoint which breaks the transaction count logic.
   config.before do |example_file|
-    ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.allow_cross_database_modification =
+    ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.suppress =
       CROSS_DB_MODIFICATION_ALLOW_LIST.include?(example_file.file_path_rerun_argument)
   end
 
+  # Reset after execution to preferred state
   config.after do |example_file|
-    ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.allow_cross_database_modification = false
+    ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.suppress = true
   end
 end