From ab952b19496c9569059f2dbb2b1dacaf44e618a8 Mon Sep 17 00:00:00 2001
From: John Hawthorn <john@hawthorn.email>
Date: Tue, 29 Jan 2019 10:05:30 -0800
Subject: Write update_last_write_timestamp after request

We need to update using the timestamp from the end of the request, not
the start. For example, if a request spends 5+ seconds writing, we still
want to wait another 5 seconds for replication lag.

Since we now run the update after we yield, we need to use ensure to
make sure we update the timestamp even if there is an exception.
---
 activerecord/test/cases/database_selector_test.rb | 37 +++++++++++++++++++++++
 1 file changed, 37 insertions(+)

(limited to 'activerecord/test')

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)
-- 
cgit v1.2.3