Merge pull request rails#54836 from kirs/with_default_isolation_level

[ActiveRecord] Introduce with_default_isolation_level

Author: byroot

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

@@ -14,12 +14,12 @@ module AbstractPool # :nodoc:
     class NullPool # :nodoc:
       include ConnectionAdapters::AbstractPool
 
-      class NullConfig # :nodoc:
+      class NullConfig
         def method_missing(...)
           nil
         end
       end
-      NULL_CONFIG = NullConfig.new # :nodoc:
+      NULL_CONFIG = NullConfig.new
 
       def initialize
         super()
@@ -48,6 +48,11 @@ def db_config
       def dirties_query_cache
         true
       end
+
+      def default_isolation_level; end
+      def default_isolation_level=(isolation_level)
+        raise NotImplementedError, "This method should never be called"
+      end
     end
 
     # = Active Record Connection Pool
@@ -706,6 +711,16 @@ def new_connection # :nodoc:
         raise ex.set_pool(self)
       end
 
+      def default_isolation_level
+        isolation_level_key = "activerecord_default_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}"
+        ActiveSupport::IsolatedExecutionState[isolation_level_key] = isolation_level
+      end
+
       private
         def connection_lease
           @leases[ActiveSupport::IsolatedExecutionState.context]

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

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

activerecord/lib/active_record/transactions.rb

@@ -234,6 +234,25 @@ def transaction(**options, &block)
         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)
+        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
+        yield
+      ensure
+        connection_pool.default_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
+      end
+
       # Returns a representation of the current transaction state,
       # which can be a top level transaction, a savepoint, or the absence of a transaction.
       #

activerecord/test/cases/transaction_isolation_test.rb

@@ -62,6 +62,32 @@ class Tag2 < ActiveRecord::Base
       assert_equal 1, Tag.count
     end
 
+    test "default_isolation_level" do
+      assert_nil Tag.default_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.transaction do
+            Tag.create!(name: "jon")
+          end
+        end
+      end
+      assert_begin_isolation_level_event(events)
+    end
+
+    test "default_isolation_level cannot be set within open transaction" do
+      assert_raises(ActiveRecord::TransactionIsolationError) do
+        Tag.transaction do
+          Tag.with_default_isolation_level(:read_committed) { }
+        end
+      end
+    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
@@ -106,5 +132,14 @@ class Tag2 < ActiveRecord::Base
         end
       end
     end
+
+    private
+      def assert_begin_isolation_level_event(events)
+        if current_adapter?(:PostgreSQLAdapter)
+          assert_equal 1, events.select { _1.match(/BEGIN ISOLATION LEVEL READ COMMITTED/) }.size
+        else
+          assert_equal 1, events.select { _1.match(/SET TRANSACTION ISOLATION LEVEL READ COMMITTED/) }.size
+        end
+      end
   end
 end