Merge pull request #38247 from rosa/save-context-with-response

Allow updating the database selector context with the response and not only the request

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)

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

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

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)