Merge pull request rails#51474 from Shopify/transaction-callback-3

Allow to register transaction callbacks outside of a record

Author: byroot

activerecord/CHANGELOG.md

@@ -1,3 +1,51 @@
+*   `ActiveRecord::Base.transaction` now yields an `ActiveRecord::Transation` object.
+
+    This allows to register callbacks on it.
+
+    ```ruby
+    Article.transaction do |transaction|
+      article.update(published: true)
+      transaction.after_commit do
+        PublishNotificationMailer.with(article: article).deliver_later
+      end
+    end
+    ```
+
+    *Jean Boussier*
+
+*   Add `ActiveRecord::Base.current_transaction`.
+
+    Returns the current transaction, to allow registering callbacks on it.
+
+    ```ruby
+    Article.current_transaction.after_commit do
+      PublishNotificationMailer.with(article: article).deliver_later
+    end
+    ```
+
+    *Jean Boussier*
+
+*   Add `ActiveRecord.after_all_transactions_commit` callback.
+
+    Useful for code that may run either inside or outside a transaction and need
+    to perform works after the state changes have been properly peristed.
+
+    ```ruby
+    def publish_article(article)
+      article.update(published: true)
+      ActiveRecord.after_all_transactions_commit do
+        PublishNotificationMailer.with(article: article).deliver_later
+      end
+    end
+    ```
+
+    In the above example, the block is either executed immediately if called outside
+    of a transaction, or called after the open transaction is committed.
+
+    If the transaction is rolled back, the block isn't called.
+
+    *Jean Boussier*
+
 *   Add the ability to ignore counter cache columns until they are backfilled
 
     Starting to use counter caches on existing large tables can be troublesome, because the column

activerecord/lib/active_record.rb

@@ -86,6 +86,7 @@ module ActiveRecord
   autoload :Timestamp
   autoload :TokenFor
   autoload :TouchLater
+  autoload :Transaction
   autoload :Transactions
   autoload :Translation
   autoload :Validations
@@ -540,6 +541,51 @@ def self.eager_load!
   def self.disconnect_all!
     ConnectionAdapters::PoolConfig.disconnect_all!
   end
+
+  # Registers a block to be called after all the current transactions have been
+  # committed.
+  #
+  # If there is no currently open transaction, the block is called immediately.
+  #
+  # If there are multiple nested transactions, the block is called after the outermost one
+  # has been committed,
+  #
+  # If any of the currently open transactions is rolled back, the block is never called.
+  #
+  # If multiple transactions are open across multiple databases, the block will be invoked
+  # if and once all of them have been committed. But note that nesting transactions across
+  # two distinct databases is a sharding anti-pattern that comes with a world of hurts.
+  def self.after_all_transactions_commit(&block)
+    open_transactions = all_open_transactions
+
+    if open_transactions.empty?
+      yield
+    elsif open_transactions.size == 1
+      open_transactions.first.after_commit(&block)
+    else
+      count = open_transactions.size
+      callback = -> do
+        count -= 1
+        block.call if count.zero?
+      end
+      open_transactions.each do |t|
+        t.after_commit(&callback)
+      end
+      open_transactions = nil # rubocop:disable Lint/UselessAssignment avoid holding it in the closure
+    end
+  end
+
+  def self.all_open_transactions # :nodoc:
+    open_transactions = []
+    Base.connection_handler.each_connection_pool do |pool|
+      if active_connection = pool.active_connection
+        if active_connection.current_transaction.open? && active_connection.current_transaction.joinable?
+          open_transactions << active_connection.current_transaction
+        end
+      end
+    end
+    open_transactions
+  end
 end
 
 ActiveSupport.on_load(:active_record) do

activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb

@@ -235,6 +235,17 @@ def truncate_tables(*table_names) # :nodoc:
       # Runs the given block in a database transaction, and returns the result
       # of the block.
       #
+      # == Transaction callbacks
+      #
+      # #transaction yields an ActiveRecord::Transaction object on which it is
+      # possible to register callback:
+      #
+      #   ActiveRecord::Base.transaction do |transaction|
+      #     transaction.before_commit { puts "before commit!" }
+      #     transaction.after_commit { puts "after commit!" }
+      #     transaction.after_rollback { puts "after rollback!" }
+      #   end
+      #
       # == Nested transactions support
       #
       # #transaction calls can be nested. By default, this makes all database
@@ -345,7 +356,7 @@ def transaction(requires_new: nil, isolation: nil, joinable: true, &block)
           if isolation
             raise ActiveRecord::TransactionIsolationError, "cannot set isolation when joining a transaction"
           end
-          yield
+          yield current_transaction
         else
           transaction_manager.within_new_transaction(isolation: isolation, joinable: joinable, &block)
         end

activerecord/lib/active_record/connection_adapters/abstract/transaction.rb

@@ -116,15 +116,19 @@ def dirty!; end
       def invalidated?; false; end
       def invalidate!; end
       def materialized?; false; end
+      def before_commit; yield; end
+      def after_commit; yield; end
+      def after_rollback; end # noop
     end
 
-    class Transaction # :nodoc:
+    class Transaction < ActiveRecord::Transaction # :nodoc:
       attr_reader :connection, :state, :savepoint_name, :isolation_level
       attr_accessor :written
 
       delegate :invalidate!, :invalidated?, to: :@state
 
       def initialize(connection, isolation: nil, joinable: true, run_commit_callbacks: false)
+        super()
         @connection = connection
         @state = TransactionState.new
         @records = nil
@@ -191,60 +195,76 @@ def restore!
       end
 
       def rollback_records
-        return unless records
-
-        ite = unique_records
+        if records
+          begin
+            ite = unique_records
 
-        instances_to_run_callbacks_on = prepare_instances_to_run_callbacks_on(ite)
+            instances_to_run_callbacks_on = prepare_instances_to_run_callbacks_on(ite)
 
-        run_action_on_records(ite, instances_to_run_callbacks_on) do |record, should_run_callbacks|
-          record.rolledback!(force_restore_state: full_rollback?, should_run_callbacks: should_run_callbacks)
-        end
-      ensure
-        ite&.each do |i|
-          i.rolledback!(force_restore_state: full_rollback?, should_run_callbacks: false)
+            run_action_on_records(ite, instances_to_run_callbacks_on) do |record, should_run_callbacks|
+              record.rolledback!(force_restore_state: full_rollback?, should_run_callbacks: should_run_callbacks)
+            end
+          ensure
+            ite&.each do |i|
+              i.rolledback!(force_restore_state: full_rollback?, should_run_callbacks: false)
+            end
+          end
         end
+
+        @callbacks&.each(&:after_rollback)
       end
 
       def before_commit_records
-        return unless records
-
         if @run_commit_callbacks
-          if ActiveRecord.before_committed_on_all_records
-            ite = unique_records
+          if records
+            if ActiveRecord.before_committed_on_all_records
+              ite = unique_records
 
-            instances_to_run_callbacks_on = records.each_with_object({}) do |record, candidates|
-              candidates[record] = record
-            end
+              instances_to_run_callbacks_on = records.each_with_object({}) do |record, candidates|
+                candidates[record] = record
+              end
 
-            run_action_on_records(ite, instances_to_run_callbacks_on) do |record, should_run_callbacks|
-              record.before_committed! if should_run_callbacks
+              run_action_on_records(ite, instances_to_run_callbacks_on) do |record, should_run_callbacks|
+                record.before_committed! if should_run_callbacks
+              end
+            else
+              records.uniq.each(&:before_committed!)
             end
-          else
-            records.uniq.each(&:before_committed!)
           end
+
+          @callbacks&.each(&:before_commit)
         end
+        # Note: When @run_commit_callbacks is false #commit_records takes care of appending
+        # remaining callbacks to the parent transaction
       end
 
       def commit_records
-        return unless records
-
-        ite = unique_records
+        if records
+          begin
+            ite = unique_records
 
-        if @run_commit_callbacks
-          instances_to_run_callbacks_on = prepare_instances_to_run_callbacks_on(ite)
+            if @run_commit_callbacks
+              instances_to_run_callbacks_on = prepare_instances_to_run_callbacks_on(ite)
 
-          run_action_on_records(ite, instances_to_run_callbacks_on) do |record, should_run_callbacks|
-            record.committed!(should_run_callbacks: should_run_callbacks)
-          end
-        else
-          while record = ite.shift
-            # if not running callbacks, only adds the record to the parent transaction
-            connection.add_transaction_record(record)
+              run_action_on_records(ite, instances_to_run_callbacks_on) do |record, should_run_callbacks|
+                record.committed!(should_run_callbacks: should_run_callbacks)
+              end
+            else
+              while record = ite.shift
+                # if not running callbacks, only adds the record to the parent transaction
+                connection.add_transaction_record(record)
+              end
+            end
+          ensure
+            ite&.each { |i| i.committed!(should_run_callbacks: false) }
           end
         end
-      ensure
-        ite&.each { |i| i.committed!(should_run_callbacks: false) }
+
+        if @run_commit_callbacks
+          @callbacks&.each(&:after_commit)
+        elsif @callbacks
+          connection.current_transaction.append_callbacks(@callbacks)
+        end
       end
 
       def full_rollback?; true; end
@@ -533,7 +553,7 @@ def within_new_transaction(isolation: nil, joinable: true)
         @connection.lock.synchronize do
           transaction = begin_transaction(isolation: isolation, joinable: joinable)
           begin
-            yield
+            yield transaction
           rescue Exception => error
             rollback_transaction
             after_failure_actions(transaction, error)
@@ -573,7 +593,7 @@ def current_transaction
       end
 
       private
-        NULL_TRANSACTION = NullTransaction.new
+        NULL_TRANSACTION = NullTransaction.new.freeze
 
         # Deallocate invalidated prepared statements outside of the transaction
         def after_failure_actions(transaction, error)

activerecord/lib/active_record/transaction.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  class Transaction
+    class Callback # :nodoc:
+      def initialize(event, callback)
+        @event = event
+        @callback = callback
+      end
+
+      def before_commit
+        @callback.call if @event == :before_commit
+      end
+
+      def after_commit
+        @callback.call if @event == :after_commit
+      end
+
+      def after_rollback
+        @callback.call if @event == :after_rollback
+      end
+    end
+
+    def initialize # :nodoc:
+      @callbacks = nil
+    end
+
+    # Registers a block to be called before the current transaction is fully committed.
+    #
+    # If there is no currently open transactions, the block is called immediately.
+    #
+    # If the current transaction has a parent transaction, the callback is transfered to
+    # the parent when the current transaction commits, or dropped when the current transaction
+    # is rolled back. This operation is repeated until the outermost transaction is reached.
+    def before_commit(&block)
+      (@callbacks ||= []) << Callback.new(:before_commit, block)
+    end
+
+    # Registers a block to be called after the current transaction is fully committed.
+    #
+    # If there is no currently open transactions, the block is called immediately.
+    #
+    # If the current transaction has a parent transaction, the callback is transfered to
+    # the parent when the current transaction commits, or dropped when the current transaction
+    # is rolled back. This operation is repeated until the outermost transaction is reached.
+    def after_commit(&block)
+      (@callbacks ||= []) << Callback.new(:after_commit, block)
+    end
+
+    # Registers a block to be called after the current transaction is rolled back.
+    #
+    # If there is no currently open transactions, the block is never called.
+    #
+    # If the current transaction is successfully committed but has a parent
+    # transaction, the callback is automatically added to the parent transaction.
+    #
+    # If the entire chain of nested transactions are all successfully committed,
+    # the block is never called.
+    def after_rollback(&block)
+      (@callbacks ||= []) << Callback.new(:after_rollback, block)
+    end
+
+    protected
+      def append_callbacks(callbacks)
+        (@callbacks ||= []).concat(callbacks)
+      end
+  end
+end

activerecord/lib/active_record/transactions.rb

@@ -224,6 +224,11 @@ def transaction(**options, &block)
         end
       end
 
+      # Returns the current transaction. See ActiveRecord::Transactions API docs.
+      def current_transaction
+        connection_pool.active_connection&.current_transaction || ConnectionAdapters::NULL_TRANSACTION
+      end
+
       def before_commit(*args, &block) # :nodoc:
         set_options_for_callbacks!(args)
         set_callback(:before_commit, :before, *args, &block)