Merge pull request #35108 from jhawthorn/db-selection-timestamp-after

Write update_last_write_timestamp after request

2 files changed, 43 insertions, 2 deletions

diff --git a/activerecord/lib/active_record/middleware/database_selector/resolver.rb b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
index 1c056e2156..930ddf900e 100644
--- a/activerecord/lib/active_record/middleware/database_selector/resolver.rb
+++ b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
@@ -64,8 +64,12 @@ module ActiveRecord
           def write_to_primary(&blk)
             ActiveRecord::Base.connected_to(role: :writing) do
               instrumenter.instrument("database_selector.active_record.wrote_to_primary") do
-                resolver.update_last_write_timestamp
-                yield
+                begin
+                  ret = yield
+                ensure
+                  resolver.update_last_write_timestamp
+                end
+                ret
               end
             end
           end
diff --git a/activerecord/test/cases/database_selector_test.rb b/activerecord/test/cases/database_selector_test.rb
index f961fa754b..6142e223ce 100644
--- a/activerecord/test/cases/database_selector_test.rb
+++ b/activerecord/test/cases/database_selector_test.rb
@@ -58,6 +58,43 @@ module ActiveRecord
       assert called
     end
 
+    def test_write_to_primary
+      resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session)
+
+      # 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]
+    end
+
+    def test_write_to_primary_with_exception
+      resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session)
+
+      # Session should start empty
+      assert_nil @session_store[:last_write]
+
+      called = false
+      assert_raises(ActiveRecord::RecordNotFound) do
+        resolver.write do
+          assert ActiveRecord::Base.connected_to?(role: :writing)
+          called = true
+          raise ActiveRecord::RecordNotFound
+        end
+      end
+      assert called
+
+      # and be populated by the last write time
+      assert @session_store[:last_write]
+    end
+
     def test_the_middleware_chooses_writing_role_with_POST_request
       middleware = ActiveRecord::Middleware::DatabaseSelector.new(lambda { |env|
         assert ActiveRecord::Base.connected_to?(role: :writing)