Move patches to patch directory, only prepend when needed

Add tests

app/models/application_record.rb

@@ -4,6 +4,7 @@ class ApplicationRecord < ActiveRecord::Base
   include DatabaseReflection
   include Transactions
   include LegacyBulkInsert
+  include CrossDatabaseModification
 
   self.abstract_class = true
 

config/initializers/transactions_track_patch.rb → app/models/concerns/cross_database_modification.rb

 # frozen_string_literal: true
 
-module CrossModificationTransactionMixin
+module CrossDatabaseModification
   extend ActiveSupport::Concern
 
+  included do
+    private_class_method :gitlab_schema
+  end
+
   class_methods do
     def transaction(**options, &block)
       super(**options) do
@@ -26,63 +30,3 @@ module CrossModificationTransactionMixin
     end
   end
 end
-
-module TransactionGitlabSchemaMixin
-  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
-
-ActiveRecord::Base.prepend(CrossModificationTransactionMixin) if Rails.env.test?
-ActiveRecord::ConnectionAdapters::RealTransaction.prepend(TransactionGitlabSchemaMixin) if Rails.env.test?
-ActiveRecord::ConnectionAdapters::SavepointTransaction.prepend(TransactionGitlabSchemaMixin) if Rails.env.test?

config/initializers/database_query_analyzers.rb

@@ -6,6 +6,9 @@ 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/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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe CrossDatabaseModification do
+  describe '.transaction' do
+    it 'adds gitlab_schema to the current transaction', :aggregate_failures do
+      recorder = ActiveRecord::QueryRecorder.new do
+        ApplicationRecord.transaction do
+          Project.first
+        end
+      end
+
+      expect(recorder.log).to include(
+        /SAVEPOINT.*gitlab_schema:gitlab_main/,
+        /SELECT.*FROM "projects"/,
+        /RELEASE SAVEPOINT.*gitlab_schema:gitlab_main/
+      )
+
+      recorder = ActiveRecord::QueryRecorder.new do
+        Ci::ApplicationRecord.transaction do
+          Project.first
+        end
+      end
+
+      expect(recorder.log).to include(
+        /SAVEPOINT.*gitlab_schema:gitlab_ci/,
+        /SELECT.*FROM "projects"/,
+        /RELEASE SAVEPOINT.*gitlab_schema:gitlab_ci/
+      )
+
+      recorder = ActiveRecord::QueryRecorder.new do
+        Project.transaction do
+          Project.first
+        end
+      end
+
+      expect(recorder.log).to include(
+        /SAVEPOINT.*gitlab_schema:gitlab_main/,
+        /SELECT.*FROM "projects"/,
+        /RELEASE SAVEPOINT.*gitlab_schema:gitlab_main/
+      )
+
+      recorder = ActiveRecord::QueryRecorder.new do
+        Ci::Pipeline.transaction do
+          Project.first
+        end
+      end
+
+      expect(recorder.log).to include(
+        /SAVEPOINT.*gitlab_schema:gitlab_ci/,
+        /SELECT.*FROM "projects"/,
+        /RELEASE SAVEPOINT.*gitlab_schema:gitlab_ci/
+      )
+    end
+
+    it 'yields' do
+      expect { |block| ApplicationRecord.transaction(&block) }.to yield_control
+    end
+  end
+end