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

rails#54836 introduced the ability to change the isolation for the
current pool. If you're changing the default transaction isolation for an
app with many dbs/shards you would need to loop through all pools and
apply the isolation. If there are hundreds of connections, that might
create a performance issue.

This change adds a new method
`ActiveRecord.with_transaction_isolation_level` which will change
the transaction isolation for any pool accessed within the block. The
method can be used in an around filter to update all transactions for an
action or middleware to switch all pools accessed in a request.

I chose to call this from `with_transaction_returning_status` and
`transaction` because called from within `connection.transaction` made
it too difficult to tell if we were changing the isolation on an open
transaction. It would either error when it wasn't actually changing or
fail to error when it was supposed to. Wrapping the calls solves this.

While implementing this I felt that the original method had a confusing
name paired with the new one so I updated it to be more explicit that
it's for only the currently current pool.

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