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.

1 files changed, 61 insertions, 6 deletions

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