Merge pull request rails#55176 from eileencodes/change-transaction-isolation-for-all-poolswithin-block

Add ability to change the transaction isolation for all pools within a block

Author: eileencodes

activerecord/CHANGELOG.md

@@ -1,3 +1,30 @@
+*   Add ability to change transaction isolation for all pools within a block.
+
+    This functionality is useful if your application needs to change the database
+    transaction isolation for a request or action.
+
+    Calling `ActiveRecord.with_transaction_isolation_level(level) {}` in an around filter or
+    middleware will set the transaction isolation for all pools accessed within the block,
+    but not for the pools that aren't.
+
+    This works with explicit and implicit transactions:
+
+    ```ruby
+    ActiveRecord.with_transaction_isolation_level(:read_committed) do
+      Tag.transaction do # opens a transaction explicitly
+        Tag.create!
+      end
+    end
+    ```
+
+    ```ruby
+    ActiveRecord.with_transaction_isolation_level(:read_committed) do
+      Tag.create! # opens a transaction implicitly
+    end
+    ```
+
+    *Eileen M. Uchitelle*
+
 *   `:class_name` is now invalid in polymorphic `belongs_to` associations.
 
     Reason is `:class_name` does not make sense in those associations because

activerecord/lib/active_record.rb

@@ -574,6 +574,23 @@ def self.all_open_transactions # :nodoc:
     end
     open_transactions
   end
+
+  def self.default_transaction_isolation_level=(isolation_level) # :nodoc:
+    ActiveSupport::IsolatedExecutionState[:active_record_transaction_isolation] = isolation_level
+  end
+
+  def self.default_transaction_isolation_level # :nodoc:
+    ActiveSupport::IsolatedExecutionState[:active_record_transaction_isolation]
+  end
+
+  # Sets a transaction isolation level for all connection pools within the block.
+  def self.with_transaction_isolation_level(isolation_level, &block)
+    original_level = self.default_transaction_isolation_level
+    self.default_transaction_isolation_level = isolation_level
+    yield
+  ensure
+    self.default_transaction_isolation_level = original_level
+  end
 end
 
 ActiveSupport.on_load(:active_record) do

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

@@ -49,8 +49,8 @@ def dirties_query_cache
         true
       end
 
-      def default_isolation_level; end
-      def default_isolation_level=(isolation_level)
+      def pool_transaction_isolation_level; end
+      def pool_transaction_isolation_level=(isolation_level)
         raise NotImplementedError, "This method should never be called"
       end
     end
@@ -428,6 +428,24 @@ def with_connection(prevent_permanent_checkout: false)
         end
       end
 
+      def with_pool_transaction_isolation_level(isolation_level, transaction_open, &block) # :nodoc:
+        if !ActiveRecord.default_transaction_isolation_level.nil?
+          begin
+            if transaction_open && self.pool_transaction_isolation_level != ActiveRecord.default_transaction_isolation_level
+              raise ActiveRecord::TransactionIsolationError, "cannot set default isolation level while transaction is open"
+            end
+
+            old_level = self.pool_transaction_isolation_level
+            self.pool_transaction_isolation_level = isolation_level
+            yield
+          ensure
+            self.pool_transaction_isolation_level = old_level
+          end
+        else
+          yield
+        end
+      end
+
       # Returns true if a connection has already been opened.
       def connected?
         synchronize { @connections.any?(&:connected?) }
@@ -711,13 +729,13 @@ def new_connection # :nodoc:
         raise ex.set_pool(self)
       end
 
-      def default_isolation_level
-        isolation_level_key = "activerecord_default_isolation_level_#{db_config.name}"
+      def pool_transaction_isolation_level
+        isolation_level_key = "activerecord_pool_transaction_isolation_level_#{db_config.name}"
         ActiveSupport::IsolatedExecutionState[isolation_level_key]
       end
 
-      def default_isolation_level=(isolation_level)
-        isolation_level_key = "activerecord_default_isolation_level_#{db_config.name}"
+      def pool_transaction_isolation_level=(isolation_level)
+        isolation_level_key = "activerecord_pool_transaction_isolation_level_#{db_config.name}"
         ActiveSupport::IsolatedExecutionState[isolation_level_key] = isolation_level
       end
 

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

@@ -620,7 +620,7 @@ def rollback_transaction(transaction = nil)
       end
 
       def within_new_transaction(isolation: nil, joinable: true)
-        isolation ||= @connection.pool.default_isolation_level
+        isolation ||= @connection.pool.pool_transaction_isolation_level
         @connection.lock.synchronize do
           transaction = begin_transaction(isolation: isolation, joinable: joinable)
           begin

activerecord/lib/active_record/transactions.rb

@@ -230,27 +230,28 @@ module ClassMethods
       # See the ConnectionAdapters::DatabaseStatements#transaction API docs.
       def transaction(**options, &block)
         with_connection do |connection|
-          connection.transaction(**options, &block)
+          connection.pool.with_pool_transaction_isolation_level(ActiveRecord.default_transaction_isolation_level, connection.transaction_open?) do
+            connection.transaction(**options, &block)
+          end
         end
       end
 
-      # Makes all transactions initiated within the block use the isolation level
-      # that you set as the default. Useful for gradually migrating apps onto new isolation level.
-      def with_default_isolation_level(isolation_level, &block)
+      # Makes all transactions the current pool use the isolation level initiated within the block.
+      def with_pool_transaction_isolation_level(isolation_level, &block)
         if current_transaction.open?
           raise ActiveRecord::TransactionIsolationError, "cannot set default isolation level while transaction is open"
         end
 
-        old_level = connection_pool.default_isolation_level
-        connection_pool.default_isolation_level = isolation_level
+        old_level = connection_pool.pool_transaction_isolation_level
+        connection_pool.pool_transaction_isolation_level = isolation_level
         yield
       ensure
-        connection_pool.default_isolation_level = old_level
+        connection_pool.pool_transaction_isolation_level = old_level
       end
 
-      # Returns the default isolation level for the connection pool, set earlier by #with_default_isolation_level.
-      def default_isolation_level
-        connection_pool.default_isolation_level
+      # Returns the default isolation level for the connection pool, set earlier by #with_pool_transaction_isolation_level.
+      def pool_transaction_isolation_level
+        connection_pool.pool_transaction_isolation_level
       end
 
       # Returns a representation of the current transaction state,
@@ -426,18 +427,20 @@ def rolledback!(force_restore_state: false, should_run_callbacks: true) # :nodoc
     # instance.
     def with_transaction_returning_status
       self.class.with_connection do |connection|
-        status = nil
-        ensure_finalize = !connection.transaction_open?
+        connection.pool.with_pool_transaction_isolation_level(ActiveRecord.default_transaction_isolation_level, connection.transaction_open?) do
+          status = nil
+          ensure_finalize = !connection.transaction_open?
 
-        connection.transaction do
-          add_to_transaction(ensure_finalize || has_transactional_callbacks?)
-          remember_transaction_record_state
+          connection.transaction do
+            add_to_transaction(ensure_finalize || has_transactional_callbacks?)
+            remember_transaction_record_state
 
-          status = yield
-          raise ActiveRecord::Rollback unless status
+            status = yield
+            raise ActiveRecord::Rollback unless status
+          end
+          @_last_transaction_return_status = status
+          status
         end
-        @_last_transaction_return_status = status
-        status
       end
     end
 

activerecord/test/cases/transaction_isolation_test.rb

@@ -29,10 +29,16 @@ class Tag2 < ActiveRecord::Base
       self.table_name = "tags"
     end
 
+    class Dog < ARUnit2Model
+      self.table_name = "dogs"
+    end
+
     setup do
       Tag.establish_connection :arunit
       Tag2.establish_connection :arunit
+      Dog.establish_connection :arunit2
       Tag.destroy_all
+      Dog.destroy_all
     end
 
     # It is impossible to properly test read uncommitted. The SQL standard only
@@ -62,16 +68,16 @@ class Tag2 < ActiveRecord::Base
       assert_equal 1, Tag.count
     end
 
-    test "default_isolation_level" do
-      assert_nil Tag.default_isolation_level
+    test "pool_transaction_isolation_level" do
+      assert_nil Tag.pool_transaction_isolation_level
 
       events = []
       ActiveSupport::Notifications.subscribed(
         -> (event) { events << event.payload[:sql] },
         "sql.active_record",
       ) do
-        Tag.with_default_isolation_level(:read_committed) do
-          assert_equal :read_committed, Tag.default_isolation_level
+        Tag.with_pool_transaction_isolation_level(:read_committed) do
+          assert_equal :read_committed, Tag.pool_transaction_isolation_level
           Tag.transaction do
             Tag.create!(name: "jon")
           end
@@ -80,24 +86,24 @@ class Tag2 < ActiveRecord::Base
       assert_begin_isolation_level_event(events)
     end
 
-    test "default_isolation_level cannot be set within open transaction" do
+    test "pool_transaction_isolation_level cannot be set within open transaction" do
       assert_raises(ActiveRecord::TransactionIsolationError) do
         Tag.transaction do
-          Tag.with_default_isolation_level(:read_committed) { }
+          Tag.with_pool_transaction_isolation_level(:read_committed) { }
         end
       end
     end
 
-    test "default_isolation_level but transaction overrides isolation" do
-      assert_nil Tag.default_isolation_level
+    test "pool_transaction_isolation_level but transaction overrides isolation" do
+      assert_nil Tag.pool_transaction_isolation_level
 
       events = []
       ActiveSupport::Notifications.subscribed(
         -> (event) { events << event.payload[:sql] },
         "sql.active_record",
       ) do
-        Tag.with_default_isolation_level(:read_committed) do
-          assert_equal :read_committed, Tag.default_isolation_level
+        Tag.with_pool_transaction_isolation_level(:read_committed) do
+          assert_equal :read_committed, Tag.pool_transaction_isolation_level
 
           Tag.transaction(isolation: :repeatable_read) do
             Tag.create!(name: "jon")
@@ -108,6 +114,99 @@ class Tag2 < ActiveRecord::Base
       assert_begin_isolation_level_event(events, isolation: "REPEATABLE READ")
     end
 
+    test "with_transaction_isolation_level explicit transaction" do
+      assert_nil ActiveRecord.default_transaction_isolation_level
+
+      events = []
+      ActiveSupport::Notifications.subscribed(
+        -> (event) { events << event.payload[:sql] },
+        "sql.active_record",
+      ) do
+        assert_nil Tag.connection_pool.pool_transaction_isolation_level
+        assert_nil Dog.connection_pool.pool_transaction_isolation_level
+
+        ActiveRecord.with_transaction_isolation_level(:read_committed) do
+          assert_equal :read_committed, ActiveRecord.default_transaction_isolation_level
+          Tag.transaction do
+            assert_equal :read_committed, Tag.connection_pool.pool_transaction_isolation_level
+            assert_equal :read_committed, Dog.connection_pool.pool_transaction_isolation_level
+
+            Tag.create!(name: "jon")
+            Dog.create!
+          end
+        end
+      end
+
+      assert_nil Tag.connection_pool.pool_transaction_isolation_level
+      assert_nil Dog.connection_pool.pool_transaction_isolation_level
+      assert_begin_isolation_level_event(events, count: 2)
+    end
+
+    test "with_transaction_isolation_level implicit transaction" do
+      assert_nil ActiveRecord.default_transaction_isolation_level
+
+      events = []
+      ActiveSupport::Notifications.subscribed(
+        -> (event) { events << event.payload[:sql] },
+        "sql.active_record",
+      ) do
+        ActiveRecord.with_transaction_isolation_level(:read_committed) do
+          assert_equal :read_committed, ActiveRecord.default_transaction_isolation_level
+
+          Tag.create!(name: "jon")
+          Dog.create!
+        end
+      end
+
+      assert_begin_isolation_level_event(events, count: 2)
+    end
+
+    test "with_transaction_isolation_level cannot be set within open transaction" do
+      Tag.transaction do
+        assert_raises(ActiveRecord::TransactionIsolationError) do
+          ActiveRecord.with_transaction_isolation_level(:repeatable_read) do
+            Tag.create!(name: "some tag")
+          end
+        end
+      end
+    end
+
+    test "with_transaction_isolation_level cannot be changed within the block" do
+      Tag.transaction do
+        assert_raises(ActiveRecord::TransactionIsolationError) do
+          ActiveRecord.with_transaction_isolation_level(:repeatable_read) do
+            Tag.transaction do
+              ActiveRecord.with_transaction_isolation_level(:serializable) do
+                assert_raises do
+                  Tag.create!(name: "some tag")
+                end
+              end
+            end
+          end
+        end
+      end
+    end
+
+    test "with_transaction_isolation_level but transaction overrides isolation" do
+      assert_nil ActiveRecord.default_transaction_isolation_level
+
+      events = []
+      ActiveSupport::Notifications.subscribed(
+        -> (event) { events << event.payload[:sql] },
+        "sql.active_record",
+      ) do
+        ActiveRecord.with_transaction_isolation_level(:read_committed) do
+          assert_equal :read_committed, ActiveRecord.default_transaction_isolation_level
+
+          Dog.transaction(isolation: :repeatable_read) do
+            Dog.create!
+          end
+        end
+      end
+
+      assert_begin_isolation_level_event(events, isolation: "REPEATABLE READ")
+    end
+
     # We are testing that a nonrepeatable read does not happen
     if ActiveRecord::Base.lease_connection.transaction_isolation_levels.include?(:repeatable_read)
       test "repeatable read" do
@@ -154,11 +253,11 @@ class Tag2 < ActiveRecord::Base
     end
 
     private
-      def assert_begin_isolation_level_event(events, isolation: "READ COMMITTED")
+      def assert_begin_isolation_level_event(events, isolation: "READ COMMITTED", count: 1)
         if current_adapter?(:PostgreSQLAdapter)
-          assert_equal 1, events.select { _1.match(/BEGIN ISOLATION LEVEL #{isolation}/) }.size
+          assert_equal count, events.select { _1.match(/BEGIN ISOLATION LEVEL #{isolation}/) }.size
         else
-          assert_equal 1, events.select { _1.match(/SET TRANSACTION ISOLATION LEVEL #{isolation}/) }.size
+          assert_equal count, events.select { _1.match(/SET TRANSACTION ISOLATION LEVEL #{isolation}/) }.size
         end
       end
   end