Add latest changes from gitlab-org/gitlab@master

1 files changed, 38 insertions, 5 deletions

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 1b25d3c2090..321890722b5 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -177,11 +177,6 @@ module Gitlab
       'unknown'
     end
 
-    # Monkeypatch rails with upgraded database observability
-    def self.install_monkey_patches
-      ActiveRecord::Base.prepend(ActiveRecordBaseTransactionMetrics)
-    end
-
     def self.read_only?
       false
     end
@@ -190,6 +185,18 @@ module Gitlab
       !read_only?
     end
 
+    # Monkeypatch rails with upgraded database observability
+    def self.install_transaction_metrics_patches!
+      ActiveRecord::Base.prepend(ActiveRecordBaseTransactionMetrics)
+    end
+
+    def self.install_transaction_context_patches!
+      ActiveRecord::ConnectionAdapters::TransactionManager
+        .prepend(TransactionManagerContext)
+      ActiveRecord::ConnectionAdapters::RealTransaction
+        .prepend(RealTransactionContext)
+    end
+
     # MonkeyPatch for ActiveRecord::Base for adding observability
     module ActiveRecordBaseTransactionMetrics
       extend ActiveSupport::Concern
@@ -204,6 +211,32 @@ module Gitlab
         end
       end
     end
+
+    # rubocop:disable Gitlab/ModuleWithInstanceVariables
+    module TransactionManagerContext
+      def transaction_context
+        @stack.first.try(:gitlab_transaction_context)
+      end
+    end
+
+    module RealTransactionContext
+      def gitlab_transaction_context
+        @gitlab_transaction_context ||= ::Gitlab::Database::Transaction::Context.new
+      end
+
+      def commit
+        gitlab_transaction_context.commit
+
+        super
+      end
+
+      def rollback
+        gitlab_transaction_context.rollback
+
+        super
+      end
+    end
+    # rubocop:enable Gitlab/ModuleWithInstanceVariables
   end
 end