Merge pull request #18936 from arthurnn/txn_callbacks

Spike on new transaction callbacks

4 files changed, 105 insertions, 4 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb b/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb
index 11440e30d4..3a1e4a4a88 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb
@@ -77,6 +77,10 @@ module ActiveRecord
         @state.set_state(:committed)
       end
 
+      def before_commit_records
+        records.uniq.each(&:before_committed!)
+      end
+
       def commit_records
         ite = records.uniq
         while record = ite.shift
@@ -159,13 +163,15 @@ module ActiveRecord
 
       def commit_transaction
         inner_transaction = @stack.pop
-        inner_transaction.commit
 
         if current_transaction.joinable?
+          inner_transaction.commit
           inner_transaction.records.each do |r|
             r.add_to_transaction
           end
         else
+          inner_transaction.before_commit_records
+          inner_transaction.commit
           inner_transaction.commit_records
         end
       end
diff --git a/activerecord/lib/active_record/core.rb b/activerecord/lib/active_record/core.rb
index de48681746..d41872d767 100644
--- a/activerecord/lib/active_record/core.rb
+++ b/activerecord/lib/active_record/core.rb
@@ -488,7 +488,7 @@ module ActiveRecord
     end
 
     def has_transactional_callbacks? # :nodoc:
-      !_rollback_callbacks.empty? || !_commit_callbacks.empty?
+      !_rollback_callbacks.empty? || !_commit_callbacks.empty? || !_before_commit_callbacks.empty?
     end
 
     # Updates the attributes on this particular ActiveRecord object so that
diff --git a/activerecord/lib/active_record/transactions.rb b/activerecord/lib/active_record/transactions.rb
index dd405c7796..598defce50 100644
--- a/activerecord/lib/active_record/transactions.rb
+++ b/activerecord/lib/active_record/transactions.rb
@@ -7,6 +7,10 @@ module ActiveRecord
 
     included do
       define_callbacks :commit, :rollback,
+                       :before_commit,
+                       :before_commit_without_transaction_enrollment,
+                       :commit_without_transaction_enrollment,
+                       :rollback_without_transaction_enrollment,
                        scope: [:kind, :name]
     end
 
@@ -206,6 +210,11 @@ module ActiveRecord
         connection.transaction(options, &block)
       end
 
+      def before_commit(*args, &block) # :nodoc:
+        set_options_for_callbacks!(args)
+        set_callback(:before_commit, :before, *args, &block)
+      end
+
       # This callback is called after a record has been created, updated, or destroyed.
       #
       # You can specify that the callback should only be fired by a certain action with
@@ -233,6 +242,21 @@ module ActiveRecord
         set_callback(:rollback, :after, *args, &block)
       end
 
+      def before_commit_without_transaction_enrollment(*args, &block) # :nodoc:
+        set_options_for_callbacks!(args)
+        set_callback(:before_commit_without_transaction_enrollment, :before, *args, &block)
+      end
+
+      def after_commit_without_transaction_enrollment(*args, &block) # :nodoc:
+        set_options_for_callbacks!(args)
+        set_callback(:commit_without_transaction_enrollment, :after, *args, &block)
+      end
+
+      def after_rollback_without_transaction_enrollment(*args, &block) # :nodoc:
+        set_options_for_callbacks!(args)
+        set_callback(:rollback_without_transaction_enrollment, :after, *args, &block)
+      end
+
       def raise_in_transactional_callbacks
         ActiveSupport::Deprecation.warn('ActiveRecord::Base.raise_in_transactional_callbacks is deprecated and will be removed without replacement.')
         true
@@ -296,12 +320,20 @@ module ActiveRecord
       clear_transaction_record_state
     end
 
+    def before_committed! # :nodoc:
+      _run_before_commit_without_transaction_enrollment_callbacks
+      _run_before_commit_callbacks
+    end
+
     # Call the +after_commit+ callbacks.
     #
     # Ensure that it is not called if the object was never persisted (failed create),
     # but call it after the commit of a destroyed object.
     def committed!(should_run_callbacks: true) #:nodoc:
-      _run_commit_callbacks if should_run_callbacks && destroyed? || persisted?
+      if should_run_callbacks && destroyed? || persisted?
+        _run_commit_without_transaction_enrollment_callbacks
+        _run_commit_callbacks
+      end
     ensure
       force_clear_transaction_record_state
     end
@@ -309,7 +341,10 @@ module ActiveRecord
     # Call the +after_rollback+ callbacks. The +force_restore_state+ argument indicates if the record
     # state should be rolled back to the beginning or just to the last savepoint.
     def rolledback!(force_restore_state: false, should_run_callbacks: true) #:nodoc:
-      _run_rollback_callbacks if should_run_callbacks
+      if should_run_callbacks
+        _run_rollback_without_transaction_enrollment_callbacks
+        _run_rollback_callbacks
+      end
     ensure
       restore_transaction_record_state(force_restore_state)
       clear_transaction_record_state
diff --git a/activerecord/test/cases/transaction_callbacks_test.rb b/activerecord/test/cases/transaction_callbacks_test.rb
index f185cda263..e868022fed 100644
--- a/activerecord/test/cases/transaction_callbacks_test.rb
+++ b/activerecord/test/cases/transaction_callbacks_test.rb
@@ -400,3 +400,63 @@ class CallbacksOnMultipleActionsTest < ActiveRecord::TestCase
     assert_equal [:update_and_destroy, :create_and_destroy], topic.history
   end
 end
+
+
+class TransactionEnrollmentCallbacksTest < ActiveRecord::TestCase
+
+  class TopicWithoutTransactionalEnrollmentCallbacks < ActiveRecord::Base
+    self.table_name = :topics
+
+    before_commit_without_transaction_enrollment { |r| r.history << :before_commit }
+    after_commit_without_transaction_enrollment { |r| r.history << :after_commit }
+    after_rollback_without_transaction_enrollment { |r| r.history << :rollback }
+
+    def history
+      @history ||= []
+    end
+  end
+
+  def setup
+    @topic = TopicWithoutTransactionalEnrollmentCallbacks.create!
+  end
+
+  def test_commit_does_not_run_transactions_callbacks_without_enrollment
+    @topic.transaction do
+      @topic.content = 'foo'
+      @topic.save!
+    end
+    assert @topic.history.empty?
+  end
+
+  def test_commit_run_transactions_callbacks_with_explicit_enrollment
+    @topic.transaction do
+      2.times do
+        @topic.content = 'foo'
+        @topic.save!
+      end
+      @topic.class.connection.add_transaction_record(@topic)
+    end
+    assert_equal [:before_commit, :after_commit], @topic.history
+  end
+
+  def test_rollback_does_not_run_transactions_callbacks_without_enrollment
+    @topic.transaction do
+      @topic.content = 'foo'
+      @topic.save!
+      raise ActiveRecord::Rollback
+    end
+    assert @topic.history.empty?
+  end
+
+  def test_rollback_run_transactions_callbacks_with_explicit_enrollment
+    @topic.transaction do
+      2.times do
+        @topic.content = 'foo'
+        @topic.save!
+      end
+      @topic.class.connection.add_transaction_record(@topic)
+      raise ActiveRecord::Rollback
+    end
+    assert_equal [:rollback], @topic.history
+  end
+end