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.

8 files changed, 111 insertions, 49 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
diff --git a/activerecord/test/cases/adapter_test.rb b/activerecord/test/cases/adapter_test.rb
index 276bdf8a8d..0bc617edbe 100644
--- a/activerecord/test/cases/adapter_test.rb
+++ b/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
diff --git a/activerecord/test/cases/adapters/mysql2/mysql2_adapter_test.rb b/activerecord/test/cases/adapters/mysql2/mysql2_adapter_test.rb
index 9364fabb35..e1f7a0b7c5 100644
--- a/activerecord/test/cases/adapters/mysql2/mysql2_adapter_test.rb
+++ b/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
diff --git a/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb b/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb
index bc517cef82..68bc87eaf8 100644
--- a/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb
+++ b/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
diff --git a/activerecord/test/cases/adapters/sqlite3/sqlite3_adapter_test.rb b/activerecord/test/cases/adapters/sqlite3/sqlite3_adapter_test.rb
index 22dc0eaad2..508f7d8945 100644
--- a/activerecord/test/cases/adapters/sqlite3/sqlite3_adapter_test.rb
+++ b/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
diff --git a/activerecord/test/cases/base_test.rb b/activerecord/test/cases/base_test.rb
index ddafa468ed..4035347d4e 100644
--- a/activerecord/test/cases/base_test.rb
+++ b/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