Merge pull request rails#52104 from Shopify/user-transaction-continued

Harden the `.current_transaction` API

Author: byroot

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

@@ -356,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 current_transaction
+          yield current_transaction.user_transaction
         else
           within_new_transaction(isolation: isolation, joinable: joinable, &block)
         end

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

@@ -277,7 +277,7 @@ def cache_notification_info(sql, name, binds)
             type_casted_binds: -> { type_casted_binds(binds) },
             name: name,
             connection: self,
-            transaction: current_transaction.presence,
+            transaction: current_transaction.user_transaction.presence,
             cached: true
           }
         end

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

@@ -108,10 +108,8 @@ def finish(outcome)
     end
 
     class NullTransaction # :nodoc:
-      def initialize; end
       def state; end
       def closed?; true; end
-      alias_method :blank?, :closed?
       def open?; false; end
       def joinable?; false; end
       def add_record(record, _ = true); end
@@ -123,12 +121,31 @@ def invalidate!; end
       def materialized?; false; end
       def before_commit; yield; end
       def after_commit; yield; end
-      def after_rollback; end # noop
-      def uuid; Digest::UUID.nil_uuid; end
+      def after_rollback; end
+      def user_transaction; ActiveRecord::Transaction::NULL_TRANSACTION; end
     end
 
-    class Transaction < ActiveRecord::Transaction # :nodoc:
-      attr_reader :connection, :state, :savepoint_name, :isolation_level
+    class Transaction # :nodoc:
+      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
+
+      attr_reader :connection, :state, :savepoint_name, :isolation_level, :user_transaction
       attr_accessor :written
 
       delegate :invalidate!, :invalidated?, to: :@state
@@ -137,14 +154,16 @@ def initialize(connection, isolation: nil, joinable: true, run_commit_callbacks:
         super()
         @connection = connection
         @state = TransactionState.new
+        @callbacks = nil
         @records = nil
         @isolation_level = isolation
         @materialized = false
         @joinable = joinable
         @run_commit_callbacks = run_commit_callbacks
         @lazy_enrollment_records = nil
         @dirty = false
-        @instrumenter = TransactionInstrumenter.new(connection: connection, transaction: self)
+        @user_transaction = joinable ? ActiveRecord::Transaction.new(self) : ActiveRecord::Transaction::NULL_TRANSACTION
+        @instrumenter = TransactionInstrumenter.new(connection: connection, transaction: @user_transaction)
       end
 
       def dirty!
@@ -155,6 +174,14 @@ def dirty?
         @dirty
       end
 
+      def open?
+        true
+      end
+
+      def closed?
+        false
+      end
+
       def add_record(record, ensure_finalize = true)
         @records ||= []
         if ensure_finalize
@@ -165,6 +192,30 @@ def add_record(record, ensure_finalize = true)
         end
       end
 
+      def before_commit(&block)
+        if @state.finalized?
+          raise ActiveRecordError, "Cannot register callbacks on a finalized transaction"
+        end
+
+        (@callbacks ||= []) << Callback.new(:before_commit, block)
+      end
+
+      def after_commit(&block)
+        if @state.finalized?
+          raise ActiveRecordError, "Cannot register callbacks on a finalized transaction"
+        end
+
+        (@callbacks ||= []) << Callback.new(:after_commit, block)
+      end
+
+      def after_rollback(&block)
+        if @state.finalized?
+          raise ActiveRecordError, "Cannot register callbacks on a finalized transaction"
+        end
+
+        (@callbacks ||= []) << Callback.new(:after_rollback, block)
+      end
+
       def records
         if @lazy_enrollment_records
           @records.concat @lazy_enrollment_records.values
@@ -276,6 +327,11 @@ def commit_records
       def full_rollback?; true; end
       def joinable?; @joinable; end
 
+      protected
+        def append_callbacks(callbacks) # :nodoc:
+          (@callbacks ||= []).concat(callbacks)
+        end
+
       private
         def unique_records
           records.uniq(&:__id__)
@@ -557,7 +613,7 @@ def within_new_transaction(isolation: nil, joinable: true)
         @connection.lock.synchronize do
           transaction = begin_transaction(isolation: isolation, joinable: joinable)
           begin
-            yield transaction
+            yield transaction.user_transaction
           rescue Exception => error
             rollback_transaction
             after_failure_actions(transaction, error)

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -1118,7 +1118,7 @@ def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name =
             statement_name:    statement_name,
             async:             async,
             connection:        self,
-            transaction:       current_transaction.presence,
+            transaction:       current_transaction.user_transaction.presence,
             row_count:         0,
             &block
           )

activerecord/lib/active_record/transaction.rb

@@ -3,9 +3,17 @@
 require "active_support/core_ext/digest"
 
 module ActiveRecord
-  # This abstract class specifies the interface to interact with the current transaction state.
+  # Class specifies the interface to interact with the current transaction state.
   #
-  # Any other methods not specified here are considered to be private interfaces.
+  # It can either map to an actual transaction or represent the abscence of a transaction.
+  #
+  # == State
+  #
+  # You can check whether a transaction is open with the +open?+ or +closed?+ methods
+  #
+  #  if Article.current_transaction.open?
+  #    # We are inside a transaction
+  #  end
   #
   # == Callbacks
   #
@@ -42,90 +50,67 @@ module ActiveRecord
   # == Caveats
   #
   # When using after_commit callbacks, it is important to note that if the callback raises an error, the transaction
-  # won't be rolled back. Relying solely on these to synchronize state between multiple systems may lead to consistency issues.
+  # won't be rolled back as it was already committed. Relying solely on these to synchronize state between multiple
+  # systems may lead to consistency issues.
   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
+    def initialize(internal_transaction) # :nodoc:
+      @internal_transaction = internal_transaction
       @uuid = 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 transferred 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.
-    #
-    # If the callback raises an error, the transaction is rolled back.
-    def before_commit(&block)
-      (@callbacks ||= []) << Callback.new(:before_commit, block)
-    end
-
-    # Registers a block to be called after the current transaction is fully committed.
+    # Registers a block to be called after the 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 transferred to
+    # If the transaction has a parent transaction, the callback is transferred 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.
     #
     # If the callback raises an error, the transaction remains committed.
+    #
+    # If the transaction is already finalized, attempting to register a callback
+    # will raise ActiveRecord::ActiveRecordError
     def after_commit(&block)
-      (@callbacks ||= []) << Callback.new(:after_commit, block)
+      if @internal_transaction.nil?
+        yield
+      else
+        @internal_transaction.after_commit(&block)
+      end
     end
 
-    # Registers a block to be called after the current transaction is rolled back.
+    # Registers a block to be called after the 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
+    # If the 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.
+    #
+    # If the transaction is already finalized, attempting to register a callback
+    # will raise ActiveRecord::ActiveRecordError
     def after_rollback(&block)
-      (@callbacks ||= []) << Callback.new(:after_rollback, block)
+      @internal_transaction&.after_rollback(&block)
     end
 
-    # Returns true if a transaction was started.
     def open?
-      true
+      @internal_transaction&.open?
     end
 
-    # Returns true if no transaction is currently active.
     def closed?
-      false
+      !open?
     end
+
     alias_method :blank?, :closed?
 
-    # Returns a UUID for this transaction.
+    # Returns a UUID for this transaction or +nil+ if no transaction is open.
     def uuid
-      @uuid ||= Digest::UUID.uuid_v4
+      if @internal_transaction
+        @uuid ||= Digest::UUID.uuid_v4
+      end
     end
 
-    protected
-      def append_callbacks(callbacks) # :nodoc:
-        (@callbacks ||= []).concat(callbacks)
-      end
+    NULL_TRANSACTION = new(nil).freeze
   end
 end

activerecord/lib/active_record/transactions.rb

@@ -243,7 +243,7 @@ def transaction(**options, &block)
       #
       # See the ActiveRecord::Transaction documentation for detailed behavior.
       def current_transaction
-        connection_pool.active_connection&.current_transaction || ConnectionAdapters::TransactionManager::NULL_TRANSACTION
+        connection_pool.active_connection&.current_transaction&.user_transaction || Transaction::NULL_TRANSACTION
       end
 
       def before_commit(*args, &block) # :nodoc:

activerecord/test/cases/instrumentation_test.rb

@@ -171,47 +171,54 @@ def test_payload_connection_with_query_cache_enabled
       ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
     end
   end
-end
-
-class ActiveRecord::TransactionInSqlActiveRecordPayloadTest < ActiveRecord::TestCase
-  # We need current_transaction to return the null transaction.
-  self.use_transactional_tests = false
 
-  def test_payload_without_an_open_transaction
-    asserted = false
+  module TransactionInSqlActiveRecordPayloadTests
+    def test_payload_without_an_open_transaction
+      asserted = false
 
-    subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |event|
-      if event.payload.fetch(:name) == "Book Count"
-        assert_nil event.payload.fetch(:transaction)
-        asserted = true
+      subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |event|
+        if event.payload.fetch(:name) == "Book Count"
+          assert_nil event.payload.fetch(:transaction)
+          asserted = true
+        end
       end
-    end
 
-    Book.count
+      Book.count
 
-    assert asserted
-  ensure
-    ActiveSupport::Notifications.unsubscribe(subscriber)
-  end
+      assert asserted
+    ensure
+      ActiveSupport::Notifications.unsubscribe(subscriber)
+    end
+
+    def test_payload_with_an_open_transaction
+      asserted = false
+      expected_transaction = nil
 
-  def test_payload_with_an_open_transaction
-    asserted = false
-    expected_transaction = nil
+      subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |event|
+        if event.payload.fetch(:name) == "Book Count"
+          assert_same expected_transaction, event.payload.fetch(:transaction)
+          asserted = true
+        end
+      end
 
-    subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |event|
-      if event.payload.fetch(:name) == "Book Count"
-        assert_same expected_transaction, event.payload.fetch(:transaction)
-        asserted = true
+      Book.transaction do |transaction|
+        expected_transaction = transaction
+        Book.count
       end
-    end
 
-    Book.transaction do |transaction|
-      expected_transaction = transaction
-      Book.count
+      assert asserted
+    ensure
+      ActiveSupport::Notifications.unsubscribe(subscriber)
     end
+  end
+
+  class TransactionInSqlActiveRecordPayloadTest < ActiveRecord::TestCase
+    include TransactionInSqlActiveRecordPayloadTests
+  end
+
+  class TransactionInSqlActiveRecordPayloadNonTransactionalTest < ActiveRecord::TestCase
+    include TransactionInSqlActiveRecordPayloadTests
 
-    assert asserted
-  ensure
-    ActiveSupport::Notifications.unsubscribe(subscriber)
+    self.use_transactional_tests = false
   end
 end