Update `QueryAnalyzers::PreventCrossDatabaseModification` to use context

Remove all duplicate processing logic and make it fit well into `rspec`.
This requires adding `allow_cross_database_modification` flag.

lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb

@@ -6,90 +6,77 @@ module Gitlab
       class PreventCrossDatabaseModification < Database::QueryAnalyzers::Base
         CrossDatabaseModificationAcrossUnsupportedTablesError = Class.new(StandardError)
 
-        # This method will allow cross database modifications within the block
-        # Example:
-        #
-        # allow_cross_database_modification_within_transaction(url: 'url-to-an-issue') do
-        #   create(:build) # inserts ci_build and project record in one transaction
-        # end
-        def self.allow_cross_database_modification_within_transaction(url:)
-          cross_database_context = self.cross_database_context
-          return yield unless cross_database_context && cross_database_context[:enabled]
-
-          transaction_tracker_enabled_was = cross_database_context[:enabled]
-          cross_database_context[:enabled] = false
+        def self.allow_cross_database_modification?
+          Thread.current[:prevent_cross_database_modification_allowed]
+        end
 
-          yield
-        ensure
-          cross_database_context[:enabled] = transaction_tracker_enabled_was if cross_database_context
+        def self.allow_cross_database_modification=(value)
+          Thread.current[:prevent_cross_database_modification_allowed] = value
         end
 
-        def self.with_cross_database_modification_prevented(log_only: false)
-          reset_cross_database_context!
-          cross_database_context.merge!(enabled: true, log_only: log_only)
+        def self.with_allow_cross_database_modification(value, &blk)
+          previous = self.allow_cross_database_modification?
+          self.allow_cross_database_modification = value
 
-          yield if block_given?
+          yield
         ensure
-          cleanup_with_cross_database_modification_prevented if block_given?
+          self.allow_cross_database_modification = previous
         end
 
-        def self.cleanup_with_cross_database_modification_prevented
-          if cross_database_context
-            cross_database_context[:enabled] = false
-          end
+        # This method will allow cross database modifications within the block
+        # Example:
+        #
+        # allow_cross_database_modification_within_transaction(url: 'url-to-an-issue') do
+        #   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)
         end
 
-        def self.cross_database_context
-          Thread.current[:transaction_tracker]
+        # 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)
         end
 
-        def self.reset_cross_database_context!
-          Thread.current[:transaction_tracker] = initial_data
-        end
+        def self.begin!
+          super
 
-        def self.initial_data
-          {
-            enabled: false,
+          context.merge!({
             transaction_depth_by_db: Hash.new { |h, k| h[k] = 0 },
-            modified_tables_by_db: Hash.new { |h, k| h[k] = Set.new },
-            log_only: false
-          }
+            modified_tables_by_db: Hash.new { |h, k| h[k] = Set.new }
+          })
         end
 
         def self.enabled?
-          true
+          ::Feature::FlipperFeature.table_exists? &&
+            Feature.enabled?(:detect_cross_database_modification, default_enabled: :yaml)
         end
 
         # rubocop:disable Metrics/AbcSize
         def self.analyze(parsed)
-          return false unless cross_database_context
-          return false unless cross_database_context[:enabled]
-
-          connection = parsed.connection
-          return false if connection.pool.instance_of?(ActiveRecord::ConnectionAdapters::NullPool)
-
+          return if self.allow_cross_database_modification?
           return if in_factory_bot_create?
 
-          database = connection.pool.db_config.name
-
+          database = ::Gitlab::Database.db_config_name(parsed.connection)
           sql = parsed.sql
 
           # We ignore BEGIN in tests as this is the outer transaction for
           # DatabaseCleaner
           if sql.start_with?('SAVEPOINT') || (!Rails.env.test? && sql.start_with?('BEGIN'))
-            cross_database_context[:transaction_depth_by_db][database] += 1
+            context[:transaction_depth_by_db][database] += 1
 
             return
           elsif sql.start_with?('RELEASE SAVEPOINT', 'ROLLBACK TO SAVEPOINT') || (!Rails.env.test? && sql.start_with?('ROLLBACK', 'COMMIT'))
-            cross_database_context[:transaction_depth_by_db][database] -= 1
-            if cross_database_context[:transaction_depth_by_db][database] <= 0
-              cross_database_context[:modified_tables_by_db][database].clear
+            context[:transaction_depth_by_db][database] -= 1
+            if context[:transaction_depth_by_db][database] <= 0
+              context[:modified_tables_by_db][database].clear
             end
 
             return
           end
 
-          return if cross_database_context[:transaction_depth_by_db].values.all?(&:zero?)
+          return if context[:transaction_depth_by_db].values.all?(&:zero?)
 
           # PgQuery might fail in some cases due to limited nesting:
           # https://github.com/pganalyze/pg_query/issues/209
@@ -107,8 +94,8 @@ module Gitlab
           # databases
           return if tables == ['schema_migrations']
 
-          cross_database_context[:modified_tables_by_db][database].merge(tables)
-          all_tables = cross_database_context[:modified_tables_by_db].values.map(&:to_a).flatten
+          context[:modified_tables_by_db][database].merge(tables)
+          all_tables = context[:modified_tables_by_db].values.map(&:to_a).flatten
           schemas = ::Gitlab::Database::GitlabSchema.table_schemas(all_tables)
 
           if schemas.many?
@@ -120,13 +107,11 @@ module Gitlab
               message += " The gitlab_schema was undefined for one or more of the tables in this transaction. Any new tables must be added to lib/gitlab/database/gitlab_schemas.yml ."
             end
 
-            begin
-              raise CrossDatabaseModificationAcrossUnsupportedTablesError, message
-            rescue CrossDatabaseModificationAcrossUnsupportedTablesError => e
-              ::Gitlab::ErrorTracking.track_exception(e, { gitlab_schemas: schemas, tables: all_tables, query: parsed.sql })
-              raise unless cross_database_context[:log_only]
-            end
+            raise CrossDatabaseModificationAcrossUnsupportedTablesError, message
           end
+        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
@@ -142,6 +127,11 @@ module Gitlab
         def self.in_factory_bot_create?
           Rails.env.test? && caller_locations.any? { |l| l.path.end_with?('lib/factory_bot/evaluation.rb') && l.label == 'create' }
         end
+
+        # When in test we raise exception
+        def self.raise_exception?
+          Rails.env.test?
+        end
       end
     end
   end

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

@@ -2,10 +2,18 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification do
+RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification, query_analyzers: false do
   let_it_be(:pipeline, refind: true) { create(:ci_pipeline) }
   let_it_be(:project, refind: true) { create(:project) }
 
+  before do
+    allow(Gitlab::Database::QueryAnalyzer.instance).to receive(:all_analyzers).and_return([described_class])
+  end
+
+  around do |example|
+    Gitlab::Database::QueryAnalyzer.instance.within { example.run }
+  end
+
   shared_examples 'successful examples' do
     context 'outside transaction' do
       it { expect { run_queries }.not_to raise_error }

spec/lib/gitlab/middleware/query_analyzer_spec.rb

@@ -21,10 +21,11 @@ RSpec.describe Gitlab::Middleware::QueryAnalyzer, query_analyzers: false do
           end
         end
 
-        it 'detects cross modifications and logs them' do
+        it 'detects cross modifications and logs them without raising exception' do
+          allow(Rails.env).to receive(:test?).and_return(false)
           expect(::Gitlab::ErrorTracking).to receive(:track_exception)
 
-          subject
+          expect { subject }.not_to raise_error
         end
 
         context 'when the detect_cross_database_modification is disabled' do

spec/lib/gitlab/sidekiq_middleware/query_analyzer_spec.rb

@@ -23,10 +23,11 @@ RSpec.describe Gitlab::SidekiqMiddleware::QueryAnalyzer, query_analyzers: false
           end
         end
 
-        it 'detects cross modifications and logs them' do
+        it 'detects cross modifications and logs them without raising exception' do
+          allow(Rails.env).to receive(:test?).and_return(false)
           expect(::Gitlab::ErrorTracking).to receive(:track_exception)
 
-          subject
+          expect { subject }.not_to raise_error
         end
 
         context 'when the detect_cross_database_modification is disabled' do

spec/support/database/prevent_cross_database_modification.rb

 # frozen_string_literal: true
 
 module PreventCrossDatabaseModificationSpecHelpers
-  def with_cross_database_modification_prevented(...)
-    ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.with_cross_database_modification_prevented(...)
-  end
-
-  def cleanup_with_cross_database_modification_prevented
-    ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.cleanup_with_cross_database_modification_prevented
-  end
+  delegate :with_cross_database_modification_prevented,
+    :allow_cross_database_modification_within_transaction,
+    to: :'::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification'
 end
 
 CROSS_DB_MODIFICATION_ALLOW_LIST = Set.new(YAML.load_file(File.join(__dir__, 'cross-database-modification-allowlist.yml'))).freeze
@@ -18,12 +14,11 @@ RSpec.configure do |config|
   # 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|
-    if CROSS_DB_MODIFICATION_ALLOW_LIST.exclude?(example_file.file_path_rerun_argument)
-      with_cross_database_modification_prevented
-    end
+    ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.allow_cross_database_modification =
+      CROSS_DB_MODIFICATION_ALLOW_LIST.include?(example_file.file_path_rerun_argument)
   end
 
   config.after do |example_file|
-    cleanup_with_cross_database_modification_prevented
+    ::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification.allow_cross_database_modification = false
   end
 end