Move while_preventing_writes from conn to handler

If we put the `while_preventing_writes` on the connection then the
middleware that sends reads to the primary and ensures they can't write
will not work. The `while_preventing_writes` will only be applied to the
connection which it's called on - which in the case of the middleware is
Ar::Base.

This worked fine if you called it directly like
`OtherDbConn.connection.while_preventing_writes` but Rails didn't have a
way of knowing you wanted to call it on all the connections.

The change here moves the `while_preventing_writes` method from the
connection to the handler so that it can block writes to all queries for
that handler. This will apply to all the connections associated with
that handler.

3 files changed, 19 insertions, 17 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
index adb5ae82a0..9b3f5260f7 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -999,15 +999,30 @@ module ActiveRecord
         end
       end
 
+      attr_reader :prevent_writes
+
       def initialize
         # These caches are keyed by spec.name (ConnectionSpecification#name).
         @owner_to_pool = ConnectionHandler.create_owner_to_pool
+        @prevent_writes = false
 
         # Backup finalizer: if the forked child never needed a pool, the above
         # early discard has not occurred
         ObjectSpace.define_finalizer self, ConnectionHandler.unowned_pool_finalizer(@owner_to_pool)
       end
 
+      # Prevent writing to the database regardless of role.
+      #
+      # In some cases you may want to prevent writes to the database
+      # even if you are on a database that can write. `while_preventing_writes`
+      # will prevent writes to the database for the duration of the block.
+      def while_preventing_writes
+        original, @prevent_writes = @prevent_writes, true
+        yield
+      ensure
+        @prevent_writes = original
+      end
+
       def connection_pool_list
         owner_to_pool.values.compact
       end
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
index 7d3ec7ecfd..c0ead17b3b 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
@@ -78,7 +78,7 @@ module ActiveRecord
       SIMPLE_INT = /\A\d+\z/
 
       attr_accessor :pool
-      attr_reader :visitor, :owner, :logger, :lock, :prepared_statements, :prevent_writes
+      attr_reader :visitor, :owner, :logger, :lock, :prepared_statements
       alias :in_use? :owner
 
       set_callback :checkin, :after, :enable_lazy_transactions!
@@ -117,7 +117,6 @@ module ActiveRecord
         @pool                = ActiveRecord::ConnectionAdapters::NullPool.new
         @idle_since          = Concurrent.monotonic_time
         @quoted_column_names, @quoted_table_names = {}, {}
-        @prevent_writes = false
         @visitor = arel_visitor
         @statements = build_statement_pool
         @lock = ActiveSupport::Concurrency::LoadInterlockAwareMonitor.new
@@ -143,19 +142,7 @@ module ActiveRecord
       # Returns true if the connection is a replica, or if +prevent_writes+
       # is set to true.
       def preventing_writes?
-        replica? || prevent_writes
-      end
-
-      # Prevent writing to the database regardless of role.
-      #
-      # In some cases you may want to prevent writes to the database
-      # even if you are on a database that can write. `while_preventing_writes`
-      # will prevent writes to the database for the duration of the block.
-      def while_preventing_writes
-        original, @prevent_writes = @prevent_writes, true
-        yield
-      ensure
-        @prevent_writes = original
+        replica? || ActiveRecord::Base.connection_handler.prevent_writes
       end
 
       def migrations_paths # :nodoc:
diff --git a/activerecord/lib/active_record/middleware/database_selector/resolver.rb b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
index 5b89c8f668..3eb1039c50 100644
--- a/activerecord/lib/active_record/middleware/database_selector/resolver.rb
+++ b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
@@ -45,8 +45,8 @@ module ActiveRecord
 
         private
           def read_from_primary(&blk)
-            ActiveRecord::Base.connection.while_preventing_writes do
-              ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role) do
+            ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role) do
+              ActiveRecord::Base.connection_handler.while_preventing_writes do
                 instrumenter.instrument("database_selector.active_record.read_from_primary") do
                   yield
                 end