Merge pull request #36843 from eileencodes/add-ability-to-unset-preve…

…nting-writes

Add ability to unset preventing writes

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -1020,8 +1020,8 @@ def initialize
       # 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

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

@@ -46,7 +46,7 @@ def write(&blk)
         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 @@ def read_from_replica(&blk)
 
           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

activerecord/test/cases/database_selector_test.rb

@@ -123,6 +123,40 @@ def test_read_from_primary_with_options
       assert read
     end
 
+    def test_preventing_writes_turns_off_for_primary_write
+      resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session, delay: 5.seconds)
+
+      # Session should start empty
+      assert_nil @session_store[:last_write]
+
+      called = false
+      resolver.write do
+        assert ActiveRecord::Base.connected_to?(role: :writing)
+        called = true
+      end
+      assert called
+
+      # and be populated by the last write time
+      assert @session_store[:last_write]
+
+      read = false
+      write = false
+      resolver.read do
+        assert ActiveRecord::Base.connected_to?(role: :writing)
+        assert ActiveRecord::Base.connection_handler.prevent_writes
+        read = true
+
+        resolver.write do
+          assert ActiveRecord::Base.connected_to?(role: :writing)
+          assert_not ActiveRecord::Base.connection_handler.prevent_writes
+          write = true
+        end
+      end
+
+      assert write
+      assert read
+    end
+
     def test_read_from_replica_with_no_delay
       resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session, delay: 0.seconds)