Merge pull request rails#44944 from ghiculescu/multi-db-reading-request

Multi database: define `reading_request?` in resolver

activerecord/CHANGELOG.md

@@ -1,3 +1,12 @@
+*   Allow overriding `reading_request?` in `DatabaseSelector::Resolver`
+
+    The default implementation checks if a request is a `get?` or `head?`,
+    but you can now change it to anything you like. If the method returns true,
+    `Resolver#read` gets called meaning the request could be served by the
+    replica database.
+
+    *Alex Ghiculescu*
+
 *   Remove `ActiveRecord.legacy_connection_handling`.
 
     *Eileen M. Uchitelle*

activerecord/lib/active_record/middleware/database_selector.rb

@@ -71,7 +71,7 @@ def select_database(request, &blk)
           context = context_klass.call(request)
           resolver = resolver_klass.call(context, options)
 
-          response = if reading_request?(request)
+          response = if resolver.reading_request?(request)
             resolver.read(&blk)
           else
             resolver.write(&blk)
@@ -80,10 +80,6 @@ def select_database(request, &blk)
           resolver.update_context(response)
           response
         end
-
-        def reading_request?(request)
-          request.get? || request.head?
-        end
     end
   end
 end

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

@@ -48,6 +48,10 @@ def update_context(response)
           context.save(response)
         end
 
+        def reading_request?(request)
+          request.get? || request.head?
+        end
+
         private
           def read_from_primary(&blk)
             ActiveRecord::Base.connected_to(role: ActiveRecord.writing_role, prevent_writes: true) do

activerecord/test/cases/database_selector_test.rb

@@ -297,5 +297,22 @@ def test_the_middleware_chooses_reading_role_with_GET_request
 
       assert_equal [200, {}, ["body"]], middleware.call("REQUEST_METHOD" => "GET")
     end
+
+    class ReadonlyResolver < ActiveRecord::Middleware::DatabaseSelector::Resolver
+      def reading_request?(request)
+        true
+      end
+    end
+
+    def test_the_middleware_chooses_reading_role_with_POST_request_if_resolver_tells_it_to
+      middleware = ActiveRecord::Middleware::DatabaseSelector.new(lambda { |env|
+        assert ActiveRecord::Base.connected_to?(role: :reading)
+        [200, {}, ["body"]]
+      }, ReadonlyResolver)
+
+      cache = ActiveSupport::Cache::MemoryStore.new
+      middleware = ActionDispatch::Session::CacheStore.new(middleware, cache: cache, key: "_session_id")
+      assert_equal [200, {}, ["body"]], middleware.call("REQUEST_METHOD" => "POST")
+    end
   end
 end