From 97c132eb467356e84571b7511149badc038d478c Mon Sep 17 00:00:00 2001
From: Adam Hegyi <ahegyi@gitlab.com>
Date: Thu, 29 Jul 2021 21:43:06 +0200
Subject: [PATCH 1/2] Utility for detecting cross DB data modification

This change adds test helpers to detect cross database data modification
within a transaction.
---
 lib/gitlab/database.rb                        |   1 +
 spec/ci_data_modification_spec.rb             | 129 ++++++++++++++++++
 spec/spec_helper.rb                           |   1 +
 spec/support/database/ci_tables.rb            |  22 +++
 .../multi_db_data_modification_check.rb       |  77 +++++++++++
 spec/support/database/prevent_cross_joins.rb  |  19 +--
 spec/tooling/quality/test_level_spec.rb       |   4 +-
 tooling/quality/test_level.rb                 |   1 +
 8 files changed, 235 insertions(+), 19 deletions(-)
 create mode 100644 spec/ci_data_modification_spec.rb
 create mode 100644 spec/support/database/ci_tables.rb
 create mode 100644 spec/support/database/multi_db_data_modification_check.rb

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 1b25d3c2090c20..522ef7c79a6844 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -198,6 +198,7 @@ module ActiveRecordBaseTransactionMetrics
         # A monkeypatch over ActiveRecord::Base.transaction.
         # It provides observability into transactional methods.
         def transaction(**options, &block)
+          ActiveSupport::Notifications.instrument('transaction.active_record.start', { connection: connection })
           ActiveSupport::Notifications.instrument('transaction.active_record', { connection: connection }) do
             super(**options, &block)
           end
diff --git a/spec/ci_data_modification_spec.rb b/spec/ci_data_modification_spec.rb
new file mode 100644
index 00000000000000..67f390358d991f
--- /dev/null
+++ b/spec/ci_data_modification_spec.rb
@@ -0,0 +1,129 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'CI data modification check in transaction', :ci_data_modification_check do
+  let_it_be(:pipeline, refind: true) { create(:ci_pipeline) }
+  let_it_be(:project, refind: true) { create(:project) }
+
+  shared_examples 'succeessful examples' do
+    context 'outside transaction' do
+      it { expect { run_queries }.not_to raise_error }
+    end
+
+    context 'within transaction' do
+      it do
+        Project.transaction do
+          expect { run_queries }.not_to raise_error
+        end
+      end
+    end
+
+    context 'within nested transaction' do
+      it do
+        Project.transaction(requires_new: true) do
+          Project.transaction(requires_new: true) do
+            expect { run_queries }.not_to raise_error
+          end
+        end
+      end
+    end
+  end
+
+  context 'when CI and other tables are read in a transaction' do
+    def run_queries
+      pipeline.reload
+      project.reload
+    end
+
+    include_examples 'succeessful examples'
+  end
+
+  context 'when only CI data is modified' do
+    def run_queries
+      pipeline.touch
+      project.reload
+    end
+
+    include_examples 'succeessful examples'
+  end
+
+  context 'when other data is modified' do
+    def run_queries
+      pipeline.reload
+      project.touch
+    end
+
+    include_examples 'succeessful examples'
+  end
+
+  context 'when both CI and other data is modified' do
+    def run_queries
+      pipeline.touch
+      project.touch
+    end
+
+    context 'outside transaction' do
+      it { expect { run_queries }.not_to raise_error }
+    end
+
+    context 'when data modification happens in a transaction' do
+      it 'raises error' do
+        Project.transaction do
+          expect { run_queries }.to raise_error /Cross-database data modification queries/
+        end
+      end
+
+      context 'when data modification happens in nested transactions' do
+        it 'raises error' do
+          Project.transaction(requires_new: true) do
+            project.touch
+            Project.transaction(requires_new: true) do
+              expect { pipeline.touch }.to raise_error /Cross-database data modification queries/
+            end
+          end
+        end
+      end
+    end
+  end
+
+  context 'when CI association is modified through project' do
+    def run_queries
+      project.variables.build(key: 'a', value: 'v')
+      project.save!
+    end
+
+    include_examples 'succeessful examples'
+  end
+
+  describe '#disable_ci_data_modification_check' do
+    it 'skips raising error' do
+      expect do
+        disable_ci_data_modification_check(url: 'gitlab-issue') do
+          Project.transaction do
+            pipeline.touch
+            project.touch
+          end
+        end
+      end.not_to raise_error
+    end
+
+    it 'raises error when complex factories are built referencing both databases' do
+      expect do
+        ApplicationRecord.transaction do
+          create(:ci_pipeline)
+        end
+      end.to raise_error /Cross-database data modification queries/
+    end
+
+    it 'skips raising error on factory creation' do
+      expect do
+        disable_ci_data_modification_check(url: 'gitlab-issue') do
+          ApplicationRecord.transaction do
+            create(:ci_pipeline)
+          end
+        end
+      end.not_to raise_error
+    end
+  end
+end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index b95b7fad5a00a8..d9e6961e3e5f2e 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -193,6 +193,7 @@
   config.include SidekiqMiddleware
   config.include StubActionCableConnection, type: :channel
   config.include StubSpamServices
+  config.include CiDataModificationCheck
 
   include StubFeatureFlags
 
diff --git a/spec/support/database/ci_tables.rb b/spec/support/database/ci_tables.rb
new file mode 100644
index 00000000000000..ee63895f081185
--- /dev/null
+++ b/spec/support/database/ci_tables.rb
@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+# This class stores the CI-related database tables which are
+# going to be moved to a separate database.
+module Database
+  module CiTables
+    def self.include?(name)
+      ci_tables.include?(name)
+    end
+
+    def self.ci_tables
+      @@ci_tables ||= Set.new.tap do |tables| # rubocop:disable Style/ClassVars
+        tables.merge(Ci::ApplicationRecord.descendants.map(&:table_name).compact)
+
+        # It was decided that taggings/tags are best placed with CI
+        # https://gitlab.com/gitlab-org/gitlab/-/issues/333413
+        tables.add('taggings')
+        tables.add('tags')
+      end
+    end
+  end
+end
diff --git a/spec/support/database/multi_db_data_modification_check.rb b/spec/support/database/multi_db_data_modification_check.rb
new file mode 100644
index 00000000000000..17cfba6243ec9e
--- /dev/null
+++ b/spec/support/database/multi_db_data_modification_check.rb
@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+module CiDataModificationCheck
+  def disable_ci_data_modification_check(url:)
+    return yield unless Thread.current[:transaction_tracker]
+
+    transaction_tracker_enabled_was = Thread.current[:transaction_tracker][:enabled]
+    Thread.current[:transaction_tracker][:enabled] = false
+
+    yield
+  ensure
+    Thread.current[:transaction_tracker][:enabled] = transaction_tracker_enabled_was
+  end
+end
+
+RSpec.configure do |config|
+  config.before(:each, :ci_data_modification_check) do
+    transaction_subscriber = ActiveSupport::Notifications.subscribe(/transaction\.active_record/) do |name, start, finish, id, payload|
+      database = payload[:connection].pool.db_config.name
+
+      if name == 'transaction.active_record.start'
+        count = Thread.current[:transaction_tracker][:open_transaction_count] += 1
+        Rails.logger.debug "-- Transaction opened on '#{database}', open transaction count: #{count}"
+      elsif name == 'transaction.active_record'
+        count = Thread.current[:transaction_tracker][:open_transaction_count] -= 1
+        Rails.logger.debug "-- Transaction closed on '#{database}', open transaction count: #{count}"
+      end
+    end
+
+    query_subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |name, start, finish, id, payload|
+      next unless Thread.current[:transaction_tracker]
+      next unless Thread.current[:transaction_tracker][:enabled]
+
+      count = Thread.current[:transaction_tracker][:open_transaction_count]
+      next if count.zero?
+
+      tables = begin
+        PgQuery.parse(payload[:sql]).dml_tables
+      rescue StandardError
+        []
+      end
+
+      next if tables.empty?
+
+      Thread.current[:transaction_tracker][:modified_tables].merge(tables)
+
+      ci_table_referenced = false
+      main_table_referenced = false
+      Thread.current[:transaction_tracker][:modified_tables].each do |table|
+        if table.start_with?('ci_') || table == 'tags' || table == 'taggings'
+          ci_table_referenced = true
+        else
+          main_table_referenced = true
+        end
+      end
+
+      if ci_table_referenced && main_table_referenced
+        message = 'Cross-database data modification queries (CI and Main) were detected within a transaction'
+        Rails.logger.debug "-- #{message}"
+        raise message
+      end
+    end
+
+    Thread.current[:transaction_tracker] = {
+      enabled: true,
+      transaction_subscriber: transaction_subscriber,
+      query_subscriber: query_subscriber,
+      open_transaction_count: 0,
+      modified_tables: Set.new
+    }
+  end
+
+  config.after(:each, :ci_data_modification_check) do
+    ActiveSupport::Notifications.unsubscribe(Thread.current[:transaction_tracker][:transaction_subscriber])
+    ActiveSupport::Notifications.unsubscribe(Thread.current[:transaction_tracker][:query_subscriber])
+  end
+end
diff --git a/spec/support/database/prevent_cross_joins.rb b/spec/support/database/prevent_cross_joins.rb
index c49b165c6c3d7f..789721ccd38a1f 100644
--- a/spec/support/database/prevent_cross_joins.rb
+++ b/spec/support/database/prevent_cross_joins.rb
@@ -37,23 +37,8 @@ def self.validate_cross_joins!(sql)
 
     # Returns true if a set includes only CI tables, or includes only non-CI tables
     def self.only_ci_or_only_main?(tables)
-      tables.all? { |table| ci_table_name?(table) } ||
-        tables.none? { |table| ci_table_name?(table) }
-    end
-
-    def self.ci_table_name?(name)
-      ci_tables.include?(name)
-    end
-
-    def self.ci_tables
-      @@ci_tables ||= Set.new.tap do |tables| # rubocop:disable Style/ClassVars
-        tables.merge(Ci::ApplicationRecord.descendants.map(&:table_name).compact)
-
-        # It was decided that taggings/tags are best placed with CI
-        # https://gitlab.com/gitlab-org/gitlab/-/issues/333413
-        tables.add('taggings')
-        tables.add('tags')
-      end
+      tables.all? { |table| CiTables.include?(table) } ||
+        tables.none? { |table| CiTables.include?(table) }
     end
 
     module SpecHelpers
diff --git a/spec/tooling/quality/test_level_spec.rb b/spec/tooling/quality/test_level_spec.rb
index 89abe337347d84..40fd3f2f0f18e4 100644
--- a/spec/tooling/quality/test_level_spec.rb
+++ b/spec/tooling/quality/test_level_spec.rb
@@ -28,7 +28,7 @@
     context 'when level is unit' do
       it 'returns a pattern' do
         expect(subject.pattern(:unit))
-          .to eq("spec/{bin,channels,config,db,dependencies,elastic,elastic_integration,experiments,factories,finders,frontend,graphql,haml_lint,helpers,initializers,javascripts,lib,models,policies,presenters,rack_servers,replicators,routing,rubocop,serializers,services,sidekiq,spam,support_specs,tasks,uploaders,validators,views,workers,tooling}{,/**/}*_spec.rb")
+          .to eq("spec/{bin,channels,ci_data_modification,config,db,dependencies,elastic,elastic_integration,experiments,factories,finders,frontend,graphql,haml_lint,helpers,initializers,javascripts,lib,models,policies,presenters,rack_servers,replicators,routing,rubocop,serializers,services,sidekiq,spam,support_specs,tasks,uploaders,validators,views,workers,tooling}{,/**/}*_spec.rb")
       end
     end
 
@@ -103,7 +103,7 @@
     context 'when level is unit' do
       it 'returns a regexp' do
         expect(subject.regexp(:unit))
-          .to eq(%r{spec/(bin|channels|config|db|dependencies|elastic|elastic_integration|experiments|factories|finders|frontend|graphql|haml_lint|helpers|initializers|javascripts|lib|models|policies|presenters|rack_servers|replicators|routing|rubocop|serializers|services|sidekiq|spam|support_specs|tasks|uploaders|validators|views|workers|tooling)})
+          .to eq(%r{spec/(bin|channels|ci_data_modification|config|db|dependencies|elastic|elastic_integration|experiments|factories|finders|frontend|graphql|haml_lint|helpers|initializers|javascripts|lib|models|policies|presenters|rack_servers|replicators|routing|rubocop|serializers|services|sidekiq|spam|support_specs|tasks|uploaders|validators|views|workers|tooling)})
       end
     end
 
diff --git a/tooling/quality/test_level.rb b/tooling/quality/test_level.rb
index ad9de067375654..34081b2b49fb7d 100644
--- a/tooling/quality/test_level.rb
+++ b/tooling/quality/test_level.rb
@@ -18,6 +18,7 @@ class TestLevel
       unit: %w[
         bin
         channels
+        ci_data_modification
         config
         db
         dependencies
-- 
GitLab


From b387d9db9fedd58e0159cfea9215093fdf17e513 Mon Sep 17 00:00:00 2001
From: Adam Hegyi <ahegyi@gitlab.com>
Date: Mon, 9 Aug 2021 07:42:27 +0200
Subject: [PATCH 2/2] Apply reviewer feedback

Apply reviewer feedback
---
 lib/gitlab/database.rb                        |   6 +-
 spec/ci_data_modification_spec.rb             | 129 ----------------
 spec/spec_helper.rb                           |   1 -
 spec/support/database/ci_tables.rb            |   2 +-
 .../multi_db_data_modification_check.rb       |  77 ----------
 .../prevent_cross_database_modification.rb    | 120 +++++++++++++++
 ...revent_cross_database_modification_spec.rb | 144 ++++++++++++++++++
 spec/tooling/quality/test_level_spec.rb       |   4 +-
 tooling/quality/test_level.rb                 |   1 -
 9 files changed, 272 insertions(+), 212 deletions(-)
 delete mode 100644 spec/ci_data_modification_spec.rb
 delete mode 100644 spec/support/database/multi_db_data_modification_check.rb
 create mode 100644 spec/support/database/prevent_cross_database_modification.rb
 create mode 100644 spec/support_specs/database/prevent_cross_database_modification_spec.rb

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 522ef7c79a6844..eb36011fcedde6 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -147,6 +147,11 @@ def self.allow_cross_joins_across_databases(url:)
       # spec/support/database/prevent_cross_joins.rb
     end
 
+    def self.allow_cross_database_modification_within_transaction(url:)
+      # this method is implemented in:
+      # spec/support/database/cross_database_modification_check.rb
+    end
+
     def self.add_post_migrate_path_to_rails(force: false)
       return if ENV['SKIP_POST_DEPLOYMENT_MIGRATIONS'] && !force
 
@@ -198,7 +203,6 @@ module ActiveRecordBaseTransactionMetrics
         # A monkeypatch over ActiveRecord::Base.transaction.
         # It provides observability into transactional methods.
         def transaction(**options, &block)
-          ActiveSupport::Notifications.instrument('transaction.active_record.start', { connection: connection })
           ActiveSupport::Notifications.instrument('transaction.active_record', { connection: connection }) do
             super(**options, &block)
           end
diff --git a/spec/ci_data_modification_spec.rb b/spec/ci_data_modification_spec.rb
deleted file mode 100644
index 67f390358d991f..00000000000000
--- a/spec/ci_data_modification_spec.rb
+++ /dev/null
@@ -1,129 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe 'CI data modification check in transaction', :ci_data_modification_check do
-  let_it_be(:pipeline, refind: true) { create(:ci_pipeline) }
-  let_it_be(:project, refind: true) { create(:project) }
-
-  shared_examples 'succeessful examples' do
-    context 'outside transaction' do
-      it { expect { run_queries }.not_to raise_error }
-    end
-
-    context 'within transaction' do
-      it do
-        Project.transaction do
-          expect { run_queries }.not_to raise_error
-        end
-      end
-    end
-
-    context 'within nested transaction' do
-      it do
-        Project.transaction(requires_new: true) do
-          Project.transaction(requires_new: true) do
-            expect { run_queries }.not_to raise_error
-          end
-        end
-      end
-    end
-  end
-
-  context 'when CI and other tables are read in a transaction' do
-    def run_queries
-      pipeline.reload
-      project.reload
-    end
-
-    include_examples 'succeessful examples'
-  end
-
-  context 'when only CI data is modified' do
-    def run_queries
-      pipeline.touch
-      project.reload
-    end
-
-    include_examples 'succeessful examples'
-  end
-
-  context 'when other data is modified' do
-    def run_queries
-      pipeline.reload
-      project.touch
-    end
-
-    include_examples 'succeessful examples'
-  end
-
-  context 'when both CI and other data is modified' do
-    def run_queries
-      pipeline.touch
-      project.touch
-    end
-
-    context 'outside transaction' do
-      it { expect { run_queries }.not_to raise_error }
-    end
-
-    context 'when data modification happens in a transaction' do
-      it 'raises error' do
-        Project.transaction do
-          expect { run_queries }.to raise_error /Cross-database data modification queries/
-        end
-      end
-
-      context 'when data modification happens in nested transactions' do
-        it 'raises error' do
-          Project.transaction(requires_new: true) do
-            project.touch
-            Project.transaction(requires_new: true) do
-              expect { pipeline.touch }.to raise_error /Cross-database data modification queries/
-            end
-          end
-        end
-      end
-    end
-  end
-
-  context 'when CI association is modified through project' do
-    def run_queries
-      project.variables.build(key: 'a', value: 'v')
-      project.save!
-    end
-
-    include_examples 'succeessful examples'
-  end
-
-  describe '#disable_ci_data_modification_check' do
-    it 'skips raising error' do
-      expect do
-        disable_ci_data_modification_check(url: 'gitlab-issue') do
-          Project.transaction do
-            pipeline.touch
-            project.touch
-          end
-        end
-      end.not_to raise_error
-    end
-
-    it 'raises error when complex factories are built referencing both databases' do
-      expect do
-        ApplicationRecord.transaction do
-          create(:ci_pipeline)
-        end
-      end.to raise_error /Cross-database data modification queries/
-    end
-
-    it 'skips raising error on factory creation' do
-      expect do
-        disable_ci_data_modification_check(url: 'gitlab-issue') do
-          ApplicationRecord.transaction do
-            create(:ci_pipeline)
-          end
-        end
-      end.not_to raise_error
-    end
-  end
-end
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index d9e6961e3e5f2e..b95b7fad5a00a8 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -193,7 +193,6 @@
   config.include SidekiqMiddleware
   config.include StubActionCableConnection, type: :channel
   config.include StubSpamServices
-  config.include CiDataModificationCheck
 
   include StubFeatureFlags
 
diff --git a/spec/support/database/ci_tables.rb b/spec/support/database/ci_tables.rb
index ee63895f081185..99fc7ac2501fcd 100644
--- a/spec/support/database/ci_tables.rb
+++ b/spec/support/database/ci_tables.rb
@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-# This class stores the CI-related database tables which are
+# This module stores the CI-related database tables which are
 # going to be moved to a separate database.
 module Database
   module CiTables
diff --git a/spec/support/database/multi_db_data_modification_check.rb b/spec/support/database/multi_db_data_modification_check.rb
deleted file mode 100644
index 17cfba6243ec9e..00000000000000
--- a/spec/support/database/multi_db_data_modification_check.rb
+++ /dev/null
@@ -1,77 +0,0 @@
-# frozen_string_literal: true
-
-module CiDataModificationCheck
-  def disable_ci_data_modification_check(url:)
-    return yield unless Thread.current[:transaction_tracker]
-
-    transaction_tracker_enabled_was = Thread.current[:transaction_tracker][:enabled]
-    Thread.current[:transaction_tracker][:enabled] = false
-
-    yield
-  ensure
-    Thread.current[:transaction_tracker][:enabled] = transaction_tracker_enabled_was
-  end
-end
-
-RSpec.configure do |config|
-  config.before(:each, :ci_data_modification_check) do
-    transaction_subscriber = ActiveSupport::Notifications.subscribe(/transaction\.active_record/) do |name, start, finish, id, payload|
-      database = payload[:connection].pool.db_config.name
-
-      if name == 'transaction.active_record.start'
-        count = Thread.current[:transaction_tracker][:open_transaction_count] += 1
-        Rails.logger.debug "-- Transaction opened on '#{database}', open transaction count: #{count}"
-      elsif name == 'transaction.active_record'
-        count = Thread.current[:transaction_tracker][:open_transaction_count] -= 1
-        Rails.logger.debug "-- Transaction closed on '#{database}', open transaction count: #{count}"
-      end
-    end
-
-    query_subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |name, start, finish, id, payload|
-      next unless Thread.current[:transaction_tracker]
-      next unless Thread.current[:transaction_tracker][:enabled]
-
-      count = Thread.current[:transaction_tracker][:open_transaction_count]
-      next if count.zero?
-
-      tables = begin
-        PgQuery.parse(payload[:sql]).dml_tables
-      rescue StandardError
-        []
-      end
-
-      next if tables.empty?
-
-      Thread.current[:transaction_tracker][:modified_tables].merge(tables)
-
-      ci_table_referenced = false
-      main_table_referenced = false
-      Thread.current[:transaction_tracker][:modified_tables].each do |table|
-        if table.start_with?('ci_') || table == 'tags' || table == 'taggings'
-          ci_table_referenced = true
-        else
-          main_table_referenced = true
-        end
-      end
-
-      if ci_table_referenced && main_table_referenced
-        message = 'Cross-database data modification queries (CI and Main) were detected within a transaction'
-        Rails.logger.debug "-- #{message}"
-        raise message
-      end
-    end
-
-    Thread.current[:transaction_tracker] = {
-      enabled: true,
-      transaction_subscriber: transaction_subscriber,
-      query_subscriber: query_subscriber,
-      open_transaction_count: 0,
-      modified_tables: Set.new
-    }
-  end
-
-  config.after(:each, :ci_data_modification_check) do
-    ActiveSupport::Notifications.unsubscribe(Thread.current[:transaction_tracker][:transaction_subscriber])
-    ActiveSupport::Notifications.unsubscribe(Thread.current[:transaction_tracker][:query_subscriber])
-  end
-end
diff --git a/spec/support/database/prevent_cross_database_modification.rb b/spec/support/database/prevent_cross_database_modification.rb
new file mode 100644
index 00000000000000..47202696dfd267
--- /dev/null
+++ b/spec/support/database/prevent_cross_database_modification.rb
@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+module Database
+  module PreventCrossDatabaseModification
+    CrossDatabaseModificationAcrossUnsupportedTablesError = Class.new(StandardError)
+
+    module GitlabDatabaseMixin
+      def allow_cross_database_modification_within_transaction(url:)
+        return yield unless Thread.current[:transaction_tracker]
+
+        cross_database_context = Database::PreventCrossDatabaseModification.cross_database_context
+        return yield unless cross_database_context[:enabled]
+
+        transaction_tracker_enabled_was = cross_database_context[:enabled]
+        cross_database_context[:enabled] = false
+
+        yield
+      ensure
+        cross_database_context[:enabled] = transaction_tracker_enabled_was if Thread.current[:transaction_tracker]
+      end
+    end
+
+    module SpecHelpers
+      def with_cross_database_modification_prevented
+        subscriber = ActiveSupport::Notifications.subscribe('sql.active_record') do |name, start, finish, id, payload|
+          PreventCrossDatabaseModification.prevent_cross_database_modification!(payload[:connection], payload[:sql])
+        end
+
+        PreventCrossDatabaseModification.reset_cross_database_context!
+        PreventCrossDatabaseModification.cross_database_context.merge!(enabled: true, subscriber: subscriber)
+
+        yield if block_given?
+      ensure
+        cleanup_with_cross_database_modification_prevented if block_given?
+      end
+
+      def cleanup_with_cross_database_modification_prevented
+        ActiveSupport::Notifications.unsubscribe(PreventCrossDatabaseModification.cross_database_context[:subscriber])
+        PreventCrossDatabaseModification.cross_database_context[:enabled] = false
+      end
+    end
+
+    def self.cross_database_context
+      Thread.current[:transaction_tracker] ||= initial_data
+    end
+
+    def self.reset_cross_database_context!
+      Thread.current[:transaction_tracker] = initial_data
+    end
+
+    def self.initial_data
+      {
+        enabled: false,
+        transaction_depth_by_db: Hash.new { |h, k| h[k] = 0 },
+        modified_tables_by_db: Hash.new { |h, k| h[k] = Set.new }
+      }
+    end
+
+    def self.prevent_cross_database_modification!(connection, sql)
+      return unless cross_database_context[:enabled]
+
+      database = connection.pool.db_config.name
+
+      if sql.start_with?('SAVEPOINT')
+        cross_database_context[:transaction_depth_by_db][database] += 1
+
+        return
+      elsif sql.start_with?('RELEASE SAVEPOINT', 'ROLLBACK TO SAVEPOINT')
+        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
+        end
+
+        return
+      end
+
+      return if cross_database_context[:transaction_depth_by_db].values.all?(&:zero?)
+
+      tables = PgQuery.parse(sql).dml_tables
+
+      return if tables.empty?
+
+      cross_database_context[:modified_tables_by_db][database].merge(tables)
+
+      ci_table_referenced = false
+      main_table_referenced = false
+      all_tables = cross_database_context[:modified_tables_by_db].values.map(&:to_a).flatten
+      all_tables.each do |table|
+        if Database::CiTables.include?(table)
+          ci_table_referenced = true
+        else
+          main_table_referenced = true
+        end
+      end
+
+      if ci_table_referenced && main_table_referenced
+        raise Database::PreventCrossDatabaseModification::CrossDatabaseModificationAcrossUnsupportedTablesError,
+          "Cross-database data modification queries (CI and Main) were detected within " \
+          "a transaction '#{all_tables.join(", ")}' discovered"
+      end
+    end
+  end
+end
+
+Gitlab::Database.singleton_class.prepend(
+  Database::PreventCrossDatabaseModification::GitlabDatabaseMixin)
+
+RSpec.configure do |config|
+  config.include(::Database::PreventCrossDatabaseModification::SpecHelpers)
+
+  # 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(:each, :prevent_cross_database_modification) do
+    with_cross_database_modification_prevented
+  end
+
+  config.after(:each, :prevent_cross_database_modification) do
+    cleanup_with_cross_database_modification_prevented
+  end
+end
diff --git a/spec/support_specs/database/prevent_cross_database_modification_spec.rb b/spec/support_specs/database/prevent_cross_database_modification_spec.rb
new file mode 100644
index 00000000000000..4fd55d59db0b7e
--- /dev/null
+++ b/spec/support_specs/database/prevent_cross_database_modification_spec.rb
@@ -0,0 +1,144 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'Database::PreventCrossDatabaseModification' do
+  let_it_be(:pipeline, refind: true) { create(:ci_pipeline) }
+  let_it_be(:project, refind: true) { create(:project) }
+
+  shared_examples 'succeessful examples' do
+    context 'outside transaction' do
+      it { expect { run_queries }.not_to raise_error }
+    end
+
+    context 'within transaction' do
+      it do
+        Project.transaction do
+          expect { run_queries }.not_to raise_error
+        end
+      end
+    end
+
+    context 'within nested transaction' do
+      it do
+        Project.transaction(requires_new: true) do
+          Project.transaction(requires_new: true) do
+            expect { run_queries }.not_to raise_error
+          end
+        end
+      end
+    end
+  end
+
+  context 'when CI and other tables are read in a transaction' do
+    def run_queries
+      pipeline.reload
+      project.reload
+    end
+
+    include_examples 'succeessful examples'
+  end
+
+  context 'when only CI data is modified' do
+    def run_queries
+      pipeline.touch
+      project.reload
+    end
+
+    include_examples 'succeessful examples'
+  end
+
+  context 'when other data is modified' do
+    def run_queries
+      pipeline.reload
+      project.touch
+    end
+
+    include_examples 'succeessful examples'
+  end
+
+  describe 'with_cross_database_modification_prevented block' do
+    it 'raises error when CI and other data is modified' do
+      expect do
+        with_cross_database_modification_prevented do
+          Project.transaction do
+            project.touch
+            pipeline.touch
+          end
+        end
+      end.to raise_error /Cross-database data modification queries/
+    end
+  end
+
+  context 'when running tests with prevent_cross_database_modification', :prevent_cross_database_modification do
+    context 'when both CI and other data is modified' do
+      def run_queries
+        project.touch
+        pipeline.touch
+      end
+
+      context 'outside transaction' do
+        it { expect { run_queries }.not_to raise_error }
+      end
+
+      context 'when data modification happens in a transaction' do
+        it 'raises error' do
+          Project.transaction do
+            expect { run_queries }.to raise_error /Cross-database data modification queries/
+          end
+        end
+
+        context 'when data modification happens in nested transactions' do
+          it 'raises error' do
+            Project.transaction(requires_new: true) do
+              project.touch
+              Project.transaction(requires_new: true) do
+                expect { pipeline.touch }.to raise_error /Cross-database data modification queries/
+              end
+            end
+          end
+        end
+      end
+    end
+
+    context 'when CI association is modified through project' do
+      def run_queries
+        project.variables.build(key: 'a', value: 'v')
+        project.save!
+      end
+
+      include_examples 'succeessful examples'
+    end
+
+    describe '#allow_cross_database_modification_within_transaction' do
+      it 'skips raising error' do
+        expect do
+          Gitlab::Database.allow_cross_database_modification_within_transaction(url: 'gitlab-issue') do
+            Project.transaction do
+              pipeline.touch
+              project.touch
+            end
+          end
+        end.not_to raise_error
+      end
+
+      it 'raises error when complex factories are built referencing both databases' do
+        expect do
+          ApplicationRecord.transaction do
+            create(:ci_pipeline)
+          end
+        end.to raise_error /Cross-database data modification queries/
+      end
+
+      it 'skips raising error on factory creation' do
+        expect do
+          Gitlab::Database.allow_cross_database_modification_within_transaction(url: 'gitlab-issue') do
+            ApplicationRecord.transaction do
+              create(:ci_pipeline)
+            end
+          end
+        end.not_to raise_error
+      end
+    end
+  end
+end
diff --git a/spec/tooling/quality/test_level_spec.rb b/spec/tooling/quality/test_level_spec.rb
index 40fd3f2f0f18e4..89abe337347d84 100644
--- a/spec/tooling/quality/test_level_spec.rb
+++ b/spec/tooling/quality/test_level_spec.rb
@@ -28,7 +28,7 @@
     context 'when level is unit' do
       it 'returns a pattern' do
         expect(subject.pattern(:unit))
-          .to eq("spec/{bin,channels,ci_data_modification,config,db,dependencies,elastic,elastic_integration,experiments,factories,finders,frontend,graphql,haml_lint,helpers,initializers,javascripts,lib,models,policies,presenters,rack_servers,replicators,routing,rubocop,serializers,services,sidekiq,spam,support_specs,tasks,uploaders,validators,views,workers,tooling}{,/**/}*_spec.rb")
+          .to eq("spec/{bin,channels,config,db,dependencies,elastic,elastic_integration,experiments,factories,finders,frontend,graphql,haml_lint,helpers,initializers,javascripts,lib,models,policies,presenters,rack_servers,replicators,routing,rubocop,serializers,services,sidekiq,spam,support_specs,tasks,uploaders,validators,views,workers,tooling}{,/**/}*_spec.rb")
       end
     end
 
@@ -103,7 +103,7 @@
     context 'when level is unit' do
       it 'returns a regexp' do
         expect(subject.regexp(:unit))
-          .to eq(%r{spec/(bin|channels|ci_data_modification|config|db|dependencies|elastic|elastic_integration|experiments|factories|finders|frontend|graphql|haml_lint|helpers|initializers|javascripts|lib|models|policies|presenters|rack_servers|replicators|routing|rubocop|serializers|services|sidekiq|spam|support_specs|tasks|uploaders|validators|views|workers|tooling)})
+          .to eq(%r{spec/(bin|channels|config|db|dependencies|elastic|elastic_integration|experiments|factories|finders|frontend|graphql|haml_lint|helpers|initializers|javascripts|lib|models|policies|presenters|rack_servers|replicators|routing|rubocop|serializers|services|sidekiq|spam|support_specs|tasks|uploaders|validators|views|workers|tooling)})
       end
     end
 
diff --git a/tooling/quality/test_level.rb b/tooling/quality/test_level.rb
index 34081b2b49fb7d..ad9de067375654 100644
--- a/tooling/quality/test_level.rb
+++ b/tooling/quality/test_level.rb
@@ -18,7 +18,6 @@ class TestLevel
       unit: %w[
         bin
         channels
-        ci_data_modification
         config
         db
         dependencies
-- 
GitLab