Merge pull request #37065 from eileencodes/push-while_preventing_writes-into-connected_to

Call `while_preventing_writes` from `connected_to`

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.
-    def connected_to(database: nil, role: nil, &blk)
+    # When using the database key a new connection will be established every time. It is not
+    # 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
+        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,14 +45,12 @@ module ActiveRecord
 
         private
           def read_from_primary(&blk)
-            ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role) do
-              ActiveRecord::Base.connection_handler.while_preventing_writes(true) do
+            ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role, prevent_writes: true) do
               instrumenter.instrument("database_selector.active_record.read_from_primary") do
                 yield
               end
             end
           end
-          end
 
           def read_from_replica(&blk)
             ActiveRecord::Base.connected_to(role: ActiveRecord::Base.reading_role) do
@@ -63,8 +61,7 @@ module ActiveRecord
           end
 
           def write_to_primary(&blk)
-            ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role) do
-              ActiveRecord::Base.connection_handler.while_preventing_writes(false) do
+            ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role, prevent_writes: false) do
               instrumenter.instrument("database_selector.active_record.wrote_to_primary") do
                 yield
               ensure
@@ -72,7 +69,6 @@ module ActiveRecord
               end
             end
           end
-          end
 
           def read_from_primary?
             !time_since_last_write_ok?

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)