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.

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

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:

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

activerecord/test/cases/adapter_test.rb

@@ -12,6 +12,7 @@ module ActiveRecord
     def setup
       @connection = ActiveRecord::Base.connection
       @connection.materialize_transactions
+      @connection_handler = ActiveRecord::Base.connection_handler
     end
 
     ##
@@ -166,7 +167,7 @@ module ActiveRecord
     def test_preventing_writes_predicate
       assert_not_predicate @connection, :preventing_writes?
 
-      @connection.while_preventing_writes do
+      @connection_handler.while_preventing_writes do
         assert_predicate @connection, :preventing_writes?
       end
 
@@ -176,7 +177,7 @@ module ActiveRecord
     def test_errors_when_an_insert_query_is_called_while_preventing_writes
       assert_no_queries do
         assert_raises(ActiveRecord::ReadOnlyError) do
-          @connection.while_preventing_writes do
+          @connection_handler.while_preventing_writes do
             @connection.transaction do
               @connection.insert("INSERT INTO subscribers(nick) VALUES ('138853948594')", nil, false)
             end
@@ -190,7 +191,7 @@ module ActiveRecord
 
       assert_no_queries do
         assert_raises(ActiveRecord::ReadOnlyError) do
-          @connection.while_preventing_writes do
+          @connection_handler.while_preventing_writes do
             @connection.transaction do
               @connection.update("UPDATE subscribers SET nick = '9989' WHERE nick = '138853948594'")
             end
@@ -204,7 +205,7 @@ module ActiveRecord
 
       assert_no_queries do
         assert_raises(ActiveRecord::ReadOnlyError) do
-          @connection.while_preventing_writes do
+          @connection_handler.while_preventing_writes do
             @connection.transaction do
               @connection.delete("DELETE FROM subscribers WHERE nick = '138853948594'")
             end
@@ -216,7 +217,7 @@ module ActiveRecord
     def test_doesnt_error_when_a_select_query_is_called_while_preventing_writes
       @connection.insert("INSERT INTO subscribers(nick) VALUES ('138853948594')")
 
-      @connection.while_preventing_writes do
+      @connection_handler.while_preventing_writes do
         result = @connection.select_all("SELECT subscribers.* FROM subscribers WHERE nick = '138853948594'")
         assert_equal 1, result.length
       end

activerecord/test/cases/adapters/mysql2/mysql2_adapter_test.rb

@@ -8,6 +8,7 @@ class Mysql2AdapterTest < ActiveRecord::Mysql2TestCase
 
   def setup
     @conn = ActiveRecord::Base.connection
+    @connection_handler = ActiveRecord::Base.connection_handler
   end
 
   def test_exec_query_nothing_raises_with_no_result_queries
@@ -148,7 +149,7 @@ class Mysql2AdapterTest < ActiveRecord::Mysql2TestCase
 
   def test_errors_when_an_insert_query_is_called_while_preventing_writes
     assert_raises(ActiveRecord::ReadOnlyError) do
-      @conn.while_preventing_writes do
+      @connection_handler.while_preventing_writes do
         @conn.insert("INSERT INTO `engines` (`car_id`) VALUES ('138853948594')")
       end
     end
@@ -158,7 +159,7 @@ class Mysql2AdapterTest < ActiveRecord::Mysql2TestCase
     @conn.insert("INSERT INTO `engines` (`car_id`) VALUES ('138853948594')")
 
     assert_raises(ActiveRecord::ReadOnlyError) do
-      @conn.while_preventing_writes do
+      @connection_handler.while_preventing_writes do
         @conn.update("UPDATE `engines` SET `engines`.`car_id` = '9989' WHERE `engines`.`car_id` = '138853948594'")
       end
     end
@@ -168,7 +169,7 @@ class Mysql2AdapterTest < ActiveRecord::Mysql2TestCase
     @conn.execute("INSERT INTO `engines` (`car_id`) VALUES ('138853948594')")
 
     assert_raises(ActiveRecord::ReadOnlyError) do
-      @conn.while_preventing_writes do
+      @connection_handler.while_preventing_writes do
         @conn.execute("DELETE FROM `engines` where `engines`.`car_id` = '138853948594'")
       end
     end
@@ -178,7 +179,7 @@ class Mysql2AdapterTest < ActiveRecord::Mysql2TestCase
     @conn.execute("INSERT INTO `engines` (`car_id`) VALUES ('138853948594')")
 
     assert_raises(ActiveRecord::ReadOnlyError) do
-      @conn.while_preventing_writes do
+      @connection_handler.while_preventing_writes do
         @conn.execute("REPLACE INTO `engines` SET `engines`.`car_id` = '249823948'")
       end
     end
@@ -187,19 +188,19 @@ class Mysql2AdapterTest < ActiveRecord::Mysql2TestCase
   def test_doesnt_error_when_a_select_query_is_called_while_preventing_writes
     @conn.execute("INSERT INTO `engines` (`car_id`) VALUES ('138853948594')")
 
-    @conn.while_preventing_writes do
+    @connection_handler.while_preventing_writes do
       assert_equal 1, @conn.execute("SELECT `engines`.* FROM `engines` WHERE `engines`.`car_id` = '138853948594'").entries.count
     end
   end
 
   def test_doesnt_error_when_a_show_query_is_called_while_preventing_writes
-    @conn.while_preventing_writes do
+    @connection_handler.while_preventing_writes do
       assert_equal 2, @conn.execute("SHOW FULL FIELDS FROM `engines`").entries.count
     end
   end
 
   def test_doesnt_error_when_a_set_query_is_called_while_preventing_writes
-    @conn.while_preventing_writes do
+    @connection_handler.while_preventing_writes do
       assert_nil @conn.execute("SET NAMES utf8mb4 COLLATE utf8mb4_unicode_ci")
     end
   end
@@ -207,7 +208,7 @@ class Mysql2AdapterTest < ActiveRecord::Mysql2TestCase
   def test_doesnt_error_when_a_read_query_with_leading_chars_is_called_while_preventing_writes
     @conn.execute("INSERT INTO `engines` (`car_id`) VALUES ('138853948594')")
 
-    @conn.while_preventing_writes do
+    @connection_handler.while_preventing_writes do
       assert_equal 1, @conn.execute("(\n( SELECT `engines`.* FROM `engines` WHERE `engines`.`car_id` = '138853948594' ) )").entries.count
     end
   end

activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb

@@ -13,6 +13,7 @@ module ActiveRecord
 
       def setup
         @connection = ActiveRecord::Base.connection
+        @connection_handler = ActiveRecord::Base.connection_handler
       end
 
       def test_bad_connection
@@ -379,7 +380,7 @@ module ActiveRecord
       def test_errors_when_an_insert_query_is_called_while_preventing_writes
         with_example_table do
           assert_raises(ActiveRecord::ReadOnlyError) do
-            @connection.while_preventing_writes do
+            @connection_handler.while_preventing_writes do
               @connection.execute("INSERT INTO ex (data) VALUES ('138853948594')")
             end
           end
@@ -391,7 +392,7 @@ module ActiveRecord
           @connection.execute("INSERT INTO ex (data) VALUES ('138853948594')")
 
           assert_raises(ActiveRecord::ReadOnlyError) do
-            @connection.while_preventing_writes do
+            @connection_handler.while_preventing_writes do
               @connection.execute("UPDATE ex SET data = '9989' WHERE data = '138853948594'")
             end
           end
@@ -403,7 +404,7 @@ module ActiveRecord
           @connection.execute("INSERT INTO ex (data) VALUES ('138853948594')")
 
           assert_raises(ActiveRecord::ReadOnlyError) do
-            @connection.while_preventing_writes do
+            @connection_handler.while_preventing_writes do
               @connection.execute("DELETE FROM ex where data = '138853948594'")
             end
           end
@@ -414,20 +415,20 @@ module ActiveRecord
         with_example_table do
           @connection.execute("INSERT INTO ex (data) VALUES ('138853948594')")
 
-          @connection.while_preventing_writes do
+          @connection_handler.while_preventing_writes do
             assert_equal 1, @connection.execute("SELECT * FROM ex WHERE data = '138853948594'").entries.count
           end
         end
       end
 
       def test_doesnt_error_when_a_show_query_is_called_while_preventing_writes
-        @connection.while_preventing_writes do
+        @connection_handler.while_preventing_writes do
           assert_equal 1, @connection.execute("SHOW TIME ZONE").entries.count
         end
       end
 
       def test_doesnt_error_when_a_set_query_is_called_while_preventing_writes
-        @connection.while_preventing_writes do
+        @connection_handler.while_preventing_writes do
           assert_equal [], @connection.execute("SET standard_conforming_strings = on").entries
         end
       end
@@ -436,7 +437,7 @@ module ActiveRecord
         with_example_table do
           @connection.execute("INSERT INTO ex (data) VALUES ('138853948594')")
 
-          @connection.while_preventing_writes do
+          @connection_handler.while_preventing_writes do
             assert_equal 1, @connection.execute("(\n( SELECT * FROM ex WHERE data = '138853948594' ) )").entries.count
           end
         end

activerecord/test/cases/adapters/sqlite3/sqlite3_adapter_test.rb

@@ -19,6 +19,8 @@ module ActiveRecord
         @conn = Base.sqlite3_connection database: ":memory:",
                                         adapter: "sqlite3",
                                         timeout: 100
+
+        @connection_handler = ActiveRecord::Base.connection_handler
       end
 
       def test_bad_connection
@@ -572,7 +574,7 @@ module ActiveRecord
       def test_errors_when_an_insert_query_is_called_while_preventing_writes
         with_example_table "id int, data string" do
           assert_raises(ActiveRecord::ReadOnlyError) do
-            @conn.while_preventing_writes do
+            @connection_handler.while_preventing_writes do
               @conn.execute("INSERT INTO ex (data) VALUES ('138853948594')")
             end
           end
@@ -584,7 +586,7 @@ module ActiveRecord
           @conn.execute("INSERT INTO ex (data) VALUES ('138853948594')")
 
           assert_raises(ActiveRecord::ReadOnlyError) do
-            @conn.while_preventing_writes do
+            @connection_handler.while_preventing_writes do
               @conn.execute("UPDATE ex SET data = '9989' WHERE data = '138853948594'")
             end
           end
@@ -596,7 +598,7 @@ module ActiveRecord
           @conn.execute("INSERT INTO ex (data) VALUES ('138853948594')")
 
           assert_raises(ActiveRecord::ReadOnlyError) do
-            @conn.while_preventing_writes do
+            @connection_handler.while_preventing_writes do
               @conn.execute("DELETE FROM ex where data = '138853948594'")
             end
           end
@@ -608,7 +610,7 @@ module ActiveRecord
           @conn.execute("INSERT INTO ex (data) VALUES ('138853948594')")
 
           assert_raises(ActiveRecord::ReadOnlyError) do
-            @conn.while_preventing_writes do
+            @connection_handler.while_preventing_writes do
               @conn.execute("REPLACE INTO ex (data) VALUES ('249823948')")
             end
           end
@@ -619,7 +621,7 @@ module ActiveRecord
         with_example_table "id int, data string" do
           @conn.execute("INSERT INTO ex (data) VALUES ('138853948594')")
 
-          @conn.while_preventing_writes do
+          @connection_handler.while_preventing_writes do
             assert_equal 1, @conn.execute("SELECT data from ex WHERE data = '138853948594'").count
           end
         end
@@ -629,7 +631,7 @@ module ActiveRecord
         with_example_table "id int, data string" do
           @conn.execute("INSERT INTO ex (data) VALUES ('138853948594')")
 
-          @conn.while_preventing_writes do
+          @connection_handler.while_preventing_writes do
             assert_equal 1, @conn.execute("  SELECT data from ex WHERE data = '138853948594'").count
           end
         end

activerecord/test/cases/base_test.rb

@@ -1496,7 +1496,7 @@ class BasicsTest < ActiveRecord::TestCase
 
   test "creating a record raises if preventing writes" do
     error = assert_raises ActiveRecord::ReadOnlyError do
-      ActiveRecord::Base.connection.while_preventing_writes do
+      ActiveRecord::Base.connection_handler.while_preventing_writes do
         Bird.create! name: "Bluejay"
       end
     end
@@ -1508,7 +1508,7 @@ class BasicsTest < ActiveRecord::TestCase
     bird = Bird.create! name: "Bluejay"
 
     error = assert_raises ActiveRecord::ReadOnlyError do
-      ActiveRecord::Base.connection.while_preventing_writes do
+      ActiveRecord::Base.connection_handler.while_preventing_writes do
         bird.update! name: "Robin"
       end
     end
@@ -1520,7 +1520,7 @@ class BasicsTest < ActiveRecord::TestCase
     bird = Bird.create! name: "Bluejay"
 
     error = assert_raises ActiveRecord::ReadOnlyError do
-      ActiveRecord::Base.connection.while_preventing_writes do
+      ActiveRecord::Base.connection_handler.while_preventing_writes do
         bird.destroy!
       end
     end
@@ -1531,7 +1531,7 @@ class BasicsTest < ActiveRecord::TestCase
   test "selecting a record does not raise if preventing writes" do
     bird = Bird.create! name: "Bluejay"
 
-    ActiveRecord::Base.connection.while_preventing_writes do
+    ActiveRecord::Base.connection_handler.while_preventing_writes do
       assert_equal bird, Bird.where(name: "Bluejay").first
     end
   end
@@ -1539,13 +1539,13 @@ class BasicsTest < ActiveRecord::TestCase
   test "an explain query does not raise if preventing writes" do
     Bird.create!(name: "Bluejay")
 
-    ActiveRecord::Base.connection.while_preventing_writes do
+    ActiveRecord::Base.connection_handler.while_preventing_writes do
       assert_queries(2) { Bird.where(name: "Bluejay").explain }
     end
   end
 
   test "an empty transaction does not raise if preventing writes" do
-    ActiveRecord::Base.connection.while_preventing_writes do
+    ActiveRecord::Base.connection_handler.while_preventing_writes do
       assert_queries(2, ignore_none: true) do
         Bird.transaction do
           ActiveRecord::Base.connection.materialize_transactions
@@ -1553,4 +1553,59 @@ class BasicsTest < ActiveRecord::TestCase
       end
     end
   end
+
+  test "preventing writes applies to all connections on a handler" do
+    conn1_error = assert_raises ActiveRecord::ReadOnlyError do
+      ActiveRecord::Base.connection_handler.while_preventing_writes do
+        assert_equal ActiveRecord::Base.connection, Bird.connection
+        assert_not_equal ARUnit2Model.connection, Bird.connection
+        Bird.create!(name: "Bluejay")
+      end
+    end
+
+    assert_match %r/\AWrite query attempted while in readonly mode: INSERT /, conn1_error.message
+
+    conn2_error = assert_raises ActiveRecord::ReadOnlyError do
+      ActiveRecord::Base.connection_handler.while_preventing_writes do
+        assert_not_equal ActiveRecord::Base.connection, Professor.connection
+        assert_equal ARUnit2Model.connection, Professor.connection
+        Professor.create!(name: "Professor Bluejay")
+      end
+    end
+
+    assert_match %r/\AWrite query attempted while in readonly mode: INSERT /, conn2_error.message
+  end
+
+  unless in_memory_db?
+    test "preventing writes with multiple handlers" do
+      ActiveRecord::Base.connects_to(database: { writing: :arunit, reading: :arunit })
+
+      conn1_error = assert_raises ActiveRecord::ReadOnlyError do
+        ActiveRecord::Base.connected_to(role: :writing) do
+          assert_equal :writing, ActiveRecord::Base.current_role
+
+          ActiveRecord::Base.connection_handler.while_preventing_writes do
+            Bird.create!(name: "Bluejay")
+          end
+        end
+      end
+
+      assert_match %r/\AWrite query attempted while in readonly mode: INSERT /, conn1_error.message
+
+      conn2_error = assert_raises ActiveRecord::ReadOnlyError do
+        ActiveRecord::Base.connected_to(role: :reading) do
+          assert_equal :reading, ActiveRecord::Base.current_role
+
+          ActiveRecord::Base.connection_handler.while_preventing_writes do
+            Bird.create!(name: "Bluejay")
+          end
+        end
+      end
+
+      assert_match %r/\AWrite query attempted while in readonly mode: INSERT /, conn2_error.message
+    ensure
+      ActiveRecord::Base.connection_handlers = { writing: ActiveRecord::Base.default_connection_handler }
+      ActiveRecord::Base.establish_connection(:arunit)
+    end
+  end
 end