Call while_preventing_writes from connected_to

If a user is using the middleware for swapping database connections and
manually calling `connected_to` in a controller/model/etc without
calling `while_preventing_writes(false)` there is potential for a race
condition where writes will be blocked.

While the user could _just_ call `while_preventing_writes` in the same
place they call `connected_to` this would mean that all cases need to
call two methods.

This PR changes `connected_to` to call `while_preventing_writes`
directly. By default we'll assume you don't want to prevent writes, but
if called with `connected_to(role: :writing, prevent_writes: true)` or
from the middleware (which calls `connected_to` this way) the writes
will be blocked.

For replicas, apps should use readonly users to enforce not writing
rather than `while_preventing_writes` directly.

Should fix the remaining issues in
https://github.com/rails/rails/issues/36830

activerecord/CHANGELOG.md

@@ -1,3 +1,11 @@
+*   Call `while_preventing_writes` directly from `connected_to`
+
+    In some cases application authors want to use the database switching middleware and make explicit calls with `connected_to. It's possible for an app to turn off writes and not turn them back on by the time we call `connected_to(role: :writing)`.
+
+    This change allows apps to fix this by assuming if a role is writing we want to allow writes, except in the case it's explicitly turned off.
+
+    *Eileen M. Uchitelle*
+
 *   Improve detection of ActiveRecord::StatementTimeout with mysql2 adapter in the edge case when the query is terminated during filesort.
 
     *Kir Shatrov*

activerecord/lib/active_record/connection_handling.rb

@@ -113,8 +113,9 @@ module ActiveRecord
     #     Dog.run_a_long_query
     #   end
     #
-    # When using the database key a new connection will be established every time.
+    # When using the database key a new connection will be established every time. It is not
-    def connected_to(database: nil, role: nil, &blk)
+    # recommended to use this outside of one-off scripts.
+    def connected_to(database: nil, role: nil, prevent_writes: false, &blk)
       if database && role
         raise ArgumentError, "connected_to can only accept a `database` or a `role` argument, but not both arguments."
       elsif database
@@ -130,7 +131,13 @@ module ActiveRecord
 
         with_handler(role, &blk)
       elsif role
-        with_handler(role.to_sym, &blk)
+        if role == writing_role
+          with_handler(role.to_sym) do
+            connection_handler.while_preventing_writes(prevent_writes, &blk)
+          end
+        else
+          with_handler(role.to_sym, &blk)
+        end
       else
         raise ArgumentError, "must provide a `database` or a `role`."
       end

activerecord/lib/active_record/middleware/database_selector/resolver.rb

@@ -45,11 +45,9 @@ module ActiveRecord
 
         private
           def read_from_primary(&blk)
-            ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role) do
+            ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role, prevent_writes: true) do
-              ActiveRecord::Base.connection_handler.while_preventing_writes(true) do
+              instrumenter.instrument("database_selector.active_record.read_from_primary") do
-                instrumenter.instrument("database_selector.active_record.read_from_primary") do
+                yield
-                  yield
-                end
               end
             end
           end
@@ -63,13 +61,11 @@ module ActiveRecord
           end
 
           def write_to_primary(&blk)
-            ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role) do
+            ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role, prevent_writes: false) do
-              ActiveRecord::Base.connection_handler.while_preventing_writes(false) do
+              instrumenter.instrument("database_selector.active_record.wrote_to_primary") do
-                instrumenter.instrument("database_selector.active_record.wrote_to_primary") do
+                yield
-                  yield
+              ensure
-                ensure
+                context.update_last_write_timestamp
-                  context.update_last_write_timestamp
-                end
               end
             end
           end

activerecord/test/cases/database_selector_test.rb

@@ -49,6 +49,24 @@ module ActiveRecord
       assert called
     end
 
+    def test_can_write_while_reading_from_replicas_if_explicit
+      @session_store[:last_write] = ActiveRecord::Middleware::DatabaseSelector::Resolver::Session.convert_time_to_timestamp(Time.now)
+
+      resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session)
+
+      called = false
+      resolver.read do
+        called = true
+        assert ActiveRecord::Base.connected_to?(role: :writing)
+        assert_predicate ActiveRecord::Base.connection, :preventing_writes?
+
+        ActiveRecord::Base.connected_to(role: :writing, prevent_writes: false) do
+          assert ActiveRecord::Base.connected_to?(role: :writing)
+        end
+      end
+      assert called
+    end
+
     def test_read_from_primary
       @session_store[:last_write] = ActiveRecord::Middleware::DatabaseSelector::Resolver::Session.convert_time_to_timestamp(Time.now)