Raise error if begin a transaction in one DB, and modify in other DB

Do not track the outer test-only transaction (joinable: false)

config/initializers/transactions_track_patch.rb

+# frozen_string_literal: true
+
+module CrossModificationTransactionMixin
+  extend ActiveSupport::Concern
+
+  class_methods do
+    def transaction(**options, &block)
+      # default value of joinable is true
+      joinable = options[:joinable].nil? || options[:joinable]
+
+      # HACK prepend_comment to get spec/lib/gitlab/database/transaction/observer_spec.rb to pass
+      marginalia_prepended = Marginalia::Comment.prepend_comment
+
+      if gitlab_schema && !marginalia_prepended && joinable
+        Marginalia.with_annotation("gitlab_schema: #{gitlab_schema}") do
+          super(**options, &block)
+        end
+      else
+        super(**options, &block)
+      end
+    end
+
+    def gitlab_schema
+      case self.name
+      when 'ActiveRecord::Base', 'ApplicationRecord'
+        :gitlab_main
+      when 'Ci::ApplicationRecord'
+        :gitlab_ci
+      else
+        Gitlab::Database::GitlabSchema.table_schema(table_name) if table_name
+      end
+    end
+  end
+end
+
+ActiveRecord::Base.prepend(CrossModificationTransactionMixin) if Rails.env.test?

lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb

@@ -27,6 +27,7 @@ module Gitlab
 
           context.merge!({
             transaction_depth_by_db: Hash.new { |h, k| h[k] = 0 },
+            transaction_schema: Hash.new { |h, k| h[k] = 0 },
             modified_tables_by_db: Hash.new { |h, k| h[k] = Set.new }
           })
         end
@@ -41,6 +42,14 @@ module Gitlab
           self.transaction_begin?(parsed) || self.transaction_end?(parsed)
         end
 
+        def self.txn_schema(sql)
+          if sql.include?('gitlab_schema: gitlab_main*/')
+            :gitlab_main
+          elsif sql.include?('gitlab_schema: gitlab_ci*/')
+            :gitlab_ci
+          end
+        end
+
         # rubocop:disable Metrics/AbcSize
         def self.analyze(parsed)
           database = ::Gitlab::Database.db_config_name(parsed.connection)
@@ -51,6 +60,9 @@ module Gitlab
           if self.transaction_begin?(parsed)
             context[:transaction_depth_by_db][database] += 1
 
+            gitlab_txn_schema = txn_schema(sql)
+            context[:transaction_schema][gitlab_txn_schema] += 1 if gitlab_txn_schema
+
             return
           elsif self.transaction_end?(parsed)
             context[:transaction_depth_by_db][database] -= 1
@@ -61,6 +73,15 @@ module Gitlab
               raise CrossDatabaseModificationAcrossUnsupportedTablesError, "Misaligned cross-DB transactions discovered at query #{sql}. This could be a bug in #{self.class} or a valid issue to investigate. Read more at https://docs.gitlab.com/ee/development/database/multiple_databases.html#removing-cross-database-transactions ."
             end
 
+            gitlab_txn_schema = txn_schema(sql)
+            if gitlab_txn_schema
+              context[:transaction_schema][gitlab_txn_schema] -= 1
+
+              if context[:transaction_schema][gitlab_txn_schema] <= 0
+                context[:transaction_schema].delete(gitlab_txn_schema)
+              end
+            end
+
             return
           end
 
@@ -87,6 +108,8 @@ module Gitlab
           all_tables = context[:modified_tables_by_db].values.map(&:to_a).flatten
           schemas = ::Gitlab::Database::GitlabSchema.table_schemas(all_tables)
 
+          schemas += context[:transaction_schema].keys
+
           if schemas.many?
             message = "Cross-database data modification of '#{schemas.to_a.join(", ")}' were detected within " \
               "a transaction modifying the '#{all_tables.to_a.join(", ")}' tables." \

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

@@ -20,7 +20,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
     end
 
     context 'within transaction' do
-      it do
+      xit do
         Project.transaction do
           expect { run_queries }.not_to raise_error
         end
@@ -28,7 +28,7 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
     end
 
     context 'within nested transaction' do
-      it do
+      xit do
         Project.transaction(requires_new: true) do
           Project.transaction(requires_new: true) do
             expect { run_queries }.not_to raise_error
@@ -54,6 +54,22 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
     end
 
     include_examples 'successful examples'
+
+    context 'in a Project transaction' do
+      it 'raises error' do
+        Project.transaction do
+          expect { run_queries }.to raise_error /Cross-database data modification/
+        end
+      end
+    end
+
+    context 'in a CI transaction' do
+      it 'does not raise error' do
+        Ci::Pipeline.transaction do
+          expect { run_queries }.not_to raise_error
+        end
+      end
+    end
   end
 
   context 'when other data is modified' do