Add ability to unset preventing writes

Previously if an app attempts to do a write inside a read request it will be
impossilbe to switch back to writing to the primary. This PR adds an
argument to the `while_preventing_writes` so that we can make sure to
turn it off if we're doing a write on a primary.

Fixes #36830

Co-authored-by: John Crepezzi <john.crepezzi@gmail.com>

2 files changed, 9 insertions, 7 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 36001efdd5..276d5a25d4 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -1020,8 +1020,8 @@ module ActiveRecord
       # 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
+      def while_preventing_writes(enabled = true)
+        original, @prevent_writes = @prevent_writes, enabled
         yield
       ensure
         @prevent_writes = original
diff --git a/activerecord/lib/active_record/middleware/database_selector/resolver.rb b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
index 3eb1039c50..b2aa453a8e 100644
--- a/activerecord/lib/active_record/middleware/database_selector/resolver.rb
+++ b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
@@ -46,7 +46,7 @@ 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 do
+              ActiveRecord::Base.connection_handler.while_preventing_writes(true) do
                 instrumenter.instrument("database_selector.active_record.read_from_primary") do
                   yield
                 end
@@ -64,10 +64,12 @@ module ActiveRecord
 
           def write_to_primary(&blk)
             ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role) do
-              instrumenter.instrument("database_selector.active_record.wrote_to_primary") do
-                yield
-              ensure
-                context.update_last_write_timestamp
+              ActiveRecord::Base.connection_handler.while_preventing_writes(false) do
+                instrumenter.instrument("database_selector.active_record.wrote_to_primary") do
+                  yield
+                ensure
+                  context.update_last_write_timestamp
+                end
               end
             end
           end