Use simpler approach as suggested by reviewer

app/models/concerns/cross_database_modification.rb

@@ -8,14 +8,18 @@ module CrossDatabaseModification
   end
 
   class_methods do
+    def gitlab_transactions_stack
+      Thread.current[:gitlab_transactions_stack] ||= []
+    end
+
     def transaction(**options, &block)
       if track_gitlab_schema_in_current_transaction?
-        super(**options) do
-          if connection.current_transaction.respond_to?(:add_gitlab_schema) && gitlab_schema
-            connection.current_transaction.add_gitlab_schema(gitlab_schema)
-          end
+        gitlab_transactions_stack.push(gitlab_schema)
 
-          yield
+        begin
+          super(**options, &block)
+        ensure
+          gitlab_transactions_stack.pop
         end
       else
         super(**options, &block)

config/initializers/0_marginalia.rb

@@ -13,7 +13,7 @@ require 'marginalia'
 # matching against the raw SQL, and prepending the comment prevents color
 # coding from working in the development log.
 Marginalia::Comment.prepend_comment = true if Rails.env.production?
-Marginalia::Comment.components = [:application, :correlation_id, :jid, :endpoint_id, :db_config_name, :gitlab_schema]
+Marginalia::Comment.components = [:application, :correlation_id, :jid, :endpoint_id, :db_config_name]
 
 # As mentioned in https://github.com/basecamp/marginalia/pull/93/files,
 # adding :line has some overhead because a regexp on the backtrace has

config/initializers/database_query_analyzers.rb

@@ -6,9 +6,6 @@ if Gitlab.dev_or_test_env? || Gitlab::Utils.to_boolean(ENV['GITLAB_ENABLE_QUERY_
   Gitlab::Database::QueryAnalyzer.instance.all_analyzers.append(::Gitlab::Database::QueryAnalyzers::GitlabSchemasMetrics)
 
   if Rails.env.test? || Gitlab::Utils.to_boolean(ENV['ENABLE_CROSS_DATABASE_MODIFICATION_DETECTION'], default: false)
-    ActiveRecord::ConnectionAdapters::RealTransaction.prepend(::Gitlab::Patch::CrossDatabaseModification::TransactionPatch)
-    ActiveRecord::ConnectionAdapters::SavepointTransaction.prepend(::Gitlab::Patch::CrossDatabaseModification::TransactionPatch)
-
     Gitlab::Database::QueryAnalyzer.instance.all_analyzers.append(::Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModification)
   end
 

lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb

@@ -27,7 +27,6 @@ module Gitlab
 
           context.merge!({
             transaction_depth_by_db: Hash.new { |h, k| h[k] = 0 },
-            transaction_schema_by_db: Hash.new { |h, k| h[k] = [] },
             modified_tables_by_db: Hash.new { |h, k| h[k] = Set.new }
           })
         end
@@ -42,14 +41,6 @@ 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)
@@ -60,25 +51,16 @@ module Gitlab
           if self.transaction_begin?(parsed)
             context[:transaction_depth_by_db][database] += 1
 
-            gitlab_txn_schema = txn_schema(sql)
-            context[:transaction_schema_by_db][database] << gitlab_txn_schema if gitlab_txn_schema
-
             return
           elsif self.transaction_end?(parsed)
             context[:transaction_depth_by_db][database] -= 1
             if context[:transaction_depth_by_db][database] == 0
               context[:modified_tables_by_db][database].clear
-              context[:transaction_schema_by_db][database].clear
             elsif context[:transaction_depth_by_db][database] < 0
               context[:transaction_depth_by_db][database] = 0
               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_by_db][database].delete(gitlab_txn_schema)
-            end
-
             return
           end
 
@@ -105,7 +87,7 @@ 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_by_db][database].uniq
+          schemas += ApplicationRecord.gitlab_transactions_stack
 
           if schemas.many?
             message = "Cross-database data modification of '#{schemas.to_a.join(", ")}' were detected within " \

lib/gitlab/marginalia/comment.rb

@@ -4,14 +4,6 @@
 module Gitlab
   module Marginalia
     module Comment
-      def self.with_gitlab_schema(gitlab_schema)
-        Thread.current[:marginalia_gitlab_schema] = gitlab_schema
-
-        yield
-      ensure
-        Thread.current[:marginalia_gitlab_schema] = nil
-      end
-
       private
 
       def jid
@@ -53,10 +45,6 @@ module Gitlab
       def db_config_name
         ::Gitlab::Database.db_config_name(marginalia_adapter)
       end
-
-      def gitlab_schema
-        Thread.current[:marginalia_gitlab_schema]
-      end
     end
   end
 end

lib/gitlab/patch/cross_database_modification/transaction_patch.rb

-# frozen_string_literal: true
-
-module Gitlab
-  module Patch
-    module CrossDatabaseModification
-      module TransactionPatch
-        extend ActiveSupport::Concern
-
-        def add_gitlab_schema(schema)
-          @gitlab_schema = schema # rubocop:disable Gitlab/ModuleWithInstanceVariables
-        end
-
-        def materialize!
-          annotate_with_gitlab_schema do
-            super
-          end
-        end
-
-        def rollback
-          annotate_with_gitlab_schema do
-            super
-          end
-        end
-
-        def commit
-          annotate_with_gitlab_schema do
-            super
-          end
-        end
-
-        private
-
-        attr_reader :gitlab_schema
-
-        # Set marginalia comment to track cross-db transactions
-        # for BEGIN/SAVEPOINT/COMMIT/RELEASE/ROLLBACK
-        def annotate_with_gitlab_schema
-          if gitlab_schema
-            Gitlab::Marginalia::Comment.with_gitlab_schema(gitlab_schema) do
-              if ENV['CROSS_DB_MOD_DEBUG']
-                debug_log(:gitlab_schema, gitlab_schema)
-
-                Gitlab::BacktraceCleaner.clean_backtrace(caller).each do |line|
-                  debug_log(:backtrace, line)
-                end
-              end
-
-              yield
-            end
-          else
-            yield
-          end
-        end
-
-        def debug_log(label, line)
-          msg = "CrossDatabaseModification #{label}:  --> #{line}"
-
-          Rails.logger.debug(msg) # rubocop:disable Gitlab/RailsLogger
-        end
-      end
-    end
-  end
-end

spec/lib/gitlab/patch/cross_database_modification/transaction_patch_spec.rb

-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Gitlab::Patch::CrossDatabaseModification::TransactionPatch do
-  let(:transaction_klass) do
-    Class.new do
-      prepend Gitlab::Patch::CrossDatabaseModification::TransactionPatch
-
-      def materialize!
-        ApplicationRecord.connection.execute(%q{SELECT 'materialize_test'})
-      end
-
-      def rollback
-        ApplicationRecord.connection.execute(%q{SELECT 'rollback_test'})
-      end
-
-      def commit
-        ApplicationRecord.connection.execute(%q{SELECT 'commit_test'})
-      end
-    end
-  end
-
-  it 'is included in Transaction classes' do
-    expect(ActiveRecord::ConnectionAdapters::RealTransaction).to include(described_class)
-    expect(ActiveRecord::ConnectionAdapters::SavepointTransaction).to include(described_class)
-  end
-
-  it 'adds a gitlab_schema comment', :aggregate_failures do
-    transaction = transaction_klass.new
-    transaction.add_gitlab_schema('_test_gitlab_schema')
-
-    recorder = ActiveRecord::QueryRecorder.new do
-      transaction.materialize!
-      transaction.commit
-    end
-
-    expect(recorder.log).to include(
-      /materialize_test.*gitlab_schema:_test_gitlab_schema/,
-      /commit_test.*gitlab_schema:_test_gitlab_schema/
-    )
-
-    recorder = ActiveRecord::QueryRecorder.new do
-      transaction.materialize!
-      transaction.rollback
-    end
-
-    expect(recorder.log).to include(
-      /materialize_test.*gitlab_schema:_test_gitlab_schema/,
-      /rollback_test.*gitlab_schema:_test_gitlab_schema/
-    )
-  end
-
-  it 'does not add a gitlab_schema comment if there is no gitlab_schema' do
-    transaction = transaction_klass.new
-
-    recorder = ActiveRecord::QueryRecorder.new do
-      transaction.materialize!
-      transaction.commit
-    end
-
-    expect(recorder.log).to include(
-      /materialize_test/,
-      /commit_test/
-    )
-
-    expect(recorder.log).not_to include(
-      /gitlab_schema/,
-      /gitlab_schema/
-    )
-  end
-
-  context 'CROSS_DB_MOD_DEBUG is enabled' do
-    before do
-      stub_env('CROSS_DB_MOD_DEBUG', '1')
-    end
-
-    it 'logs to Rails log' do
-      transaction = transaction_klass.new
-      transaction.add_gitlab_schema('_test_gitlab_schema')
-
-      allow(Rails.logger).to receive(:debug).and_call_original
-      expect(Rails.logger).to receive(:debug).with(/CrossDatabaseModification gitlab_schema:  --> _test_gitlab_schema/).and_call_original
-
-      transaction.materialize!
-      transaction.commit
-    end
-  end
-end

spec/models/concerns/cross_database_modification_spec.rb

@@ -4,31 +4,19 @@ require 'spec_helper'
 
 RSpec.describe CrossDatabaseModification do
   describe '.transaction' do
-    def assert_no_gitlab_schema_comment(log)
-      include_arg = [/gitlab_schema/] * log.size
-
-      expect(log).not_to include(*include_arg)
-    end
-
     context 'feature flag disabled' do
       before do
         stub_feature_flags(track_gitlab_schema_in_current_transaction: false)
       end
 
-      it 'does not add gitlab_schema comment' do
-        recorder = ActiveRecord::QueryRecorder.new do
-          ApplicationRecord.transaction do
-            Project.first
-          end
-        end
+      it 'does not add to gitlab_transactions_stack' do
+        ApplicationRecord.transaction do
+          expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
 
-        expect(recorder.log).to include(
-          /SAVEPOINT/,
-          /SELECT.*FROM "projects"/,
-          /RELEASE SAVEPOINT/
-        )
+          Project.first
+        end
 
-        assert_no_gitlab_schema_comment(recorder.log)
+        expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
       end
     end
 
@@ -37,71 +25,63 @@ RSpec.describe CrossDatabaseModification do
         allow(Feature::FlipperFeature).to receive(:table_exists?).and_raise(ActiveRecord::NoDatabaseError)
       end
 
-      it 'does not add gitlab_schema comment' do
-        recorder = ActiveRecord::QueryRecorder.new do
-          ApplicationRecord.transaction do
-            Project.first
-          end
-        end
+      it 'does not add to gitlab_transactions_stack' do
+        ApplicationRecord.transaction do
+          expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
 
-        expect(recorder.log).to include(
-          /SAVEPOINT/,
-          /SELECT.*FROM "projects"/,
-          /RELEASE SAVEPOINT/
-        )
+          Project.first
+        end
 
-        assert_no_gitlab_schema_comment(recorder.log)
+        expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
       end
     end
 
-    it 'adds gitlab_schema to the current transaction', :aggregate_failures do
-      recorder = ActiveRecord::QueryRecorder.new do
-        ApplicationRecord.transaction do
-          Project.first
-        end
+    it 'adds the current gitlab schema to gitlab_transactions_stack', :aggregate_failures do
+      ApplicationRecord.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_main)
+
+        Project.first
       end
 
-      expect(recorder.log).to include(
-        /SAVEPOINT.*gitlab_schema:gitlab_main/,
-        /SELECT.*FROM "projects"/,
-        /RELEASE SAVEPOINT.*gitlab_schema:gitlab_main/
-      )
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
 
       recorder = ActiveRecord::QueryRecorder.new do
         Ci::ApplicationRecord.transaction do
+          expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_ci)
+
           Project.first
         end
       end
 
-      expect(recorder.log).to include(
-        /SAVEPOINT.*gitlab_schema:gitlab_ci/,
-        /SELECT.*FROM "projects"/,
-        /RELEASE SAVEPOINT.*gitlab_schema:gitlab_ci/
-      )
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
 
-      recorder = ActiveRecord::QueryRecorder.new do
-        Project.transaction do
-          Project.first
-        end
+      Project.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_main)
+
+        Project.first
       end
 
-      expect(recorder.log).to include(
-        /SAVEPOINT.*gitlab_schema:gitlab_main/,
-        /SELECT.*FROM "projects"/,
-        /RELEASE SAVEPOINT.*gitlab_schema:gitlab_main/
-      )
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+
+      Ci::Pipeline.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_ci)
+
+        Project.first
+      end
+
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
+
+      ApplicationRecord.transaction do
+        expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_main)
 
-      recorder = ActiveRecord::QueryRecorder.new do
         Ci::Pipeline.transaction do
+          expect(ApplicationRecord.gitlab_transactions_stack).to contain_exactly(:gitlab_main, :gitlab_ci)
+
           Project.first
         end
       end
 
-      expect(recorder.log).to include(
-        /SAVEPOINT.*gitlab_schema:gitlab_ci/,
-        /SELECT.*FROM "projects"/,
-        /RELEASE SAVEPOINT.*gitlab_schema:gitlab_ci/
-      )
+      expect(ApplicationRecord.gitlab_transactions_stack).to be_empty
     end
 
     it 'yields' do