From 62e05ce96f771f2a292ad1ce87faaac1a1d886c7 Mon Sep 17 00:00:00 2001 From: Rosa Gutierrez Date: Thu, 16 Jan 2020 19:25:22 +0100 Subject: [PATCH] Allow updating the database selector context with the response Currently the database selector middleware only passes the request to the context. This is enough for the resolver to decide whether to switch to the primary, but for a custom resolver and a custom context class, it's not enough to persist any information for subsequent requests. For example, say you want to use a cookie to decide whether when to switch. It's not possible to set the response cookie from this middleware, since neither the context nor the resolver have access to it. This includes an extra step to update the context after the response has been computed. --- .../middleware/database_selector.rb | 5 ++- .../middleware/database_selector/resolver.rb | 4 +++ .../database_selector/resolver/session.rb | 3 ++ .../test/cases/database_selector_test.rb | 32 +++++++++++++++++++ 4 files changed, 43 insertions(+), 1 deletion(-) diff --git a/activerecord/lib/active_record/middleware/database_selector.rb b/activerecord/lib/active_record/middleware/database_selector.rb index 7374107048..3538d203ff 100644 --- a/activerecord/lib/active_record/middleware/database_selector.rb +++ b/activerecord/lib/active_record/middleware/database_selector.rb @@ -59,11 +59,14 @@ module ActiveRecord context = context_klass.call(request) resolver = resolver_klass.call(context, options) - if reading_request?(request) + response = if reading_request?(request) resolver.read(&blk) else resolver.write(&blk) end + + resolver.update_context(response) + response end def reading_request?(request) diff --git a/activerecord/lib/active_record/middleware/database_selector/resolver.rb b/activerecord/lib/active_record/middleware/database_selector/resolver.rb index b19700a5f7..1c98d7d1eb 100644 --- a/activerecord/lib/active_record/middleware/database_selector/resolver.rb +++ b/activerecord/lib/active_record/middleware/database_selector/resolver.rb @@ -43,6 +43,10 @@ module ActiveRecord write_to_primary(&blk) end + def update_context(response) + context.save(response) + end + private def read_from_primary(&blk) ActiveRecord::Base.connected_to(role: ActiveRecord::Base.writing_role, prevent_writes: true) do diff --git a/activerecord/lib/active_record/middleware/database_selector/resolver/session.rb b/activerecord/lib/active_record/middleware/database_selector/resolver/session.rb index df7af054b7..530701fb8d 100644 --- a/activerecord/lib/active_record/middleware/database_selector/resolver/session.rb +++ b/activerecord/lib/active_record/middleware/database_selector/resolver/session.rb @@ -38,6 +38,9 @@ module ActiveRecord def update_last_write_timestamp session[:last_write] = self.class.convert_time_to_timestamp(Time.now) end + + def save(response) + end end end end diff --git a/activerecord/test/cases/database_selector_test.rb b/activerecord/test/cases/database_selector_test.rb index 07257a9f4f..9742d34aef 100644 --- a/activerecord/test/cases/database_selector_test.rb +++ b/activerecord/test/cases/database_selector_test.rb @@ -106,6 +106,38 @@ module ActiveRecord assert @session_store[:last_write] end + def test_write_to_primary_and_update_custom_context + custom_context = Class.new(ActiveRecord::Middleware::DatabaseSelector::Resolver::Session) do + def update_last_write_timestamp + super + @wrote_to_primary = true + end + + def save(response) + response[:wrote_to_primary] = @wrote_to_primary + end + end + + resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(custom_context.new(@session_store)) + + # 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 + response = {} + resolver.update_context(response) + + # and be populated by the last write time + assert @session_store[:last_write] + # plus the response updated + assert response[:wrote_to_primary] + end + def test_write_to_primary_with_exception resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session)