Use one marginalia comment instead of two

Also add debug tooling for determining where a transaction came from

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]
+Marginalia::Comment.components = [:application, :correlation_id, :jid, :endpoint_id, :db_config_name, :gitlab_schema]
 
 # 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/transactions_track_patch.rb

@@ -5,18 +5,12 @@ module CrossModificationTransactionMixin
 
   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)
+      super(**options) do
+        if connection.current_transaction.respond_to?(:add_gitlab_schema) && gitlab_schema
+          connection.current_transaction.add_gitlab_schema(gitlab_schema)
         end
-      else
-        super(**options, &block)
+
+        yield
       end
     end
 
@@ -33,4 +27,62 @@ module CrossModificationTransactionMixin
   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?

lib/gitlab/database/query_analyzers/prevent_cross_database_modification.rb

@@ -27,7 +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 },
+            transaction_schema_by_db: Hash.new { |h, k| h[k] = [] },
             modified_tables_by_db: Hash.new { |h, k| h[k] = Set.new }
           })
         end
@@ -43,9 +43,9 @@ module Gitlab
         end
 
         def self.txn_schema(sql)
-          if sql.include?('gitlab_schema: gitlab_main*/')
+          if sql.include?('gitlab_schema:gitlab_main*/')
             :gitlab_main
-          elsif sql.include?('gitlab_schema: gitlab_ci*/')
+          elsif sql.include?('gitlab_schema:gitlab_ci*/')
             :gitlab_ci
           end
         end
@@ -61,13 +61,14 @@ module Gitlab
             context[:transaction_depth_by_db][database] += 1
 
             gitlab_txn_schema = txn_schema(sql)
-            context[:transaction_schema][gitlab_txn_schema] += 1 if gitlab_txn_schema
+            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 ."
@@ -75,11 +76,7 @@ module Gitlab
 
             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
+              context[:transaction_schema_by_db][database].delete(gitlab_txn_schema)
             end
 
             return
@@ -108,7 +105,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].keys
+          schemas += context[:transaction_schema_by_db][database].uniq
 
           if schemas.many?
             message = "Cross-database data modification of '#{schemas.to_a.join(", ")}' were detected within " \

lib/gitlab/marginalia/comment.rb

@@ -4,6 +4,14 @@
 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
@@ -45,6 +53,10 @@ 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

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

@@ -14,23 +14,25 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
     Gitlab::Database::QueryAnalyzer.instance.within { example.run }
   end
 
-  shared_examples 'successful examples' do
+  shared_examples 'successful examples' do |model:|
+    let(:model) { model }
+
     context 'outside transaction' do
       it { expect { run_queries }.not_to raise_error }
     end
 
-    context 'within transaction' do
-      xit do
-        Project.transaction do
+    context "within #{model} transaction" do
+      it do
+        model.transaction do
           expect { run_queries }.not_to raise_error
         end
       end
     end
 
-    context 'within nested transaction' do
-      xit do
-        Project.transaction(requires_new: true) do
-          Project.transaction(requires_new: true) do
+    context "within nested #{model} transaction" do
+      it do
+        model.transaction(requires_new: true) do
+          model.transaction(requires_new: true) do
             expect { run_queries }.not_to raise_error
           end
         end
@@ -38,13 +40,26 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
     end
   end
 
+  shared_examples 'cross-database modification errors' do |model:|
+    let(:model) { model }
+
+    context "within #{model} transaction" do
+      it 'raises error' do
+        model.transaction do
+          expect { run_queries }.to raise_error /Cross-database data modification/
+        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 'successful examples'
+    include_examples 'successful examples', model: Project
+    include_examples 'successful examples', model: Ci::Pipeline
   end
 
   context 'when only CI data is modified' do
@@ -53,23 +68,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
       project.reload
     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
+    include_examples 'successful examples', model: Ci::Pipeline
 
-    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
+    include_examples 'cross-database modification errors', model: Project
   end
 
   context 'when other data is modified' do
@@ -78,7 +79,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
       project.touch
     end
 
-    include_examples 'successful examples'
+    include_examples 'successful examples', model: Project
+
+    include_examples 'cross-database modification errors', model: Ci::Pipeline
   end
 
   context 'when both CI and other data is modified' do
@@ -160,7 +163,9 @@ RSpec.describe Gitlab::Database::QueryAnalyzers::PreventCrossDatabaseModificatio
         project.save!
       end
 
-      include_examples 'successful examples'
+      include_examples 'successful examples', model: Ci::Pipeline
+
+      include_examples 'cross-database modification errors', model: Project
     end
 
     describe '.allow_cross_database_modification_within_transaction' do