Merge pull request #35130 from rails/move-delay-to-options-argument

Refactor options for database selector middleware

4 files changed, 61 insertions, 11 deletions

diff --git a/activerecord/lib/active_record/middleware/database_selector.rb b/activerecord/lib/active_record/middleware/database_selector.rb
index adcfca4f8d..3ab50f5f6b 100644
--- a/activerecord/lib/active_record/middleware/database_selector.rb
+++ b/activerecord/lib/active_record/middleware/database_selector.rb
@@ -35,13 +35,14 @@ module ActiveRecord
     #   config.active_record.database_resolver = MyResolver
     #   config.active_record.database_operations = MyResolver::MySession
     class DatabaseSelector
-      def initialize(app, resolver_klass = Resolver, operations_klass = Resolver::Session)
+      def initialize(app, resolver_klass = Resolver, operations_klass = Resolver::Session, options = {})
         @app = app
         @resolver_klass = resolver_klass
         @operations_klass = operations_klass
+        @options = options
       end
 
-      attr_reader :resolver_klass, :operations_klass
+      attr_reader :resolver_klass, :operations_klass, :options
 
       # Middleware that determines which database connection to use in a multiple
       # database application.
@@ -57,7 +58,7 @@ module ActiveRecord
 
         def select_database(request, &blk)
           operations = operations_klass.build(request)
-          database_resolver = resolver_klass.call(operations)
+          database_resolver = resolver_klass.call(operations, options)
 
           if reading_request?(request)
             database_resolver.read(&blk)
diff --git a/activerecord/lib/active_record/middleware/database_selector/resolver.rb b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
index acdb9b3238..0eeb0453ef 100644
--- a/activerecord/lib/active_record/middleware/database_selector/resolver.rb
+++ b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
@@ -18,16 +18,18 @@ module ActiveRecord
       class Resolver # :nodoc:
         SEND_TO_REPLICA_DELAY = 2.seconds
 
-        def self.call(resolver)
-          new(resolver)
+        def self.call(resolver, options = {})
+          new(resolver, options)
         end
 
-        def initialize(resolver)
+        def initialize(resolver, options = {})
           @resolver = resolver
+          @options = options
+          @delay = @options && @options[:delay] ? @options[:delay] : SEND_TO_REPLICA_DELAY
           @instrumenter = ActiveSupport::Notifications.instrumenter
         end
 
-        attr_reader :resolver, :instrumenter
+        attr_reader :resolver, :delay, :instrumenter
 
         def read(&blk)
           if read_from_primary?
@@ -76,8 +78,7 @@ module ActiveRecord
           end
 
           def send_to_replica_delay
-            (ActiveRecord::Base.database_selector && ActiveRecord::Base.database_selector[:delay]) ||
-              SEND_TO_REPLICA_DELAY
+            delay
           end
 
           def time_since_last_write_ok?
diff --git a/activerecord/lib/active_record/railtie.rb b/activerecord/lib/active_record/railtie.rb
index a981fa97d9..017e279080 100644
--- a/activerecord/lib/active_record/railtie.rb
+++ b/activerecord/lib/active_record/railtie.rb
@@ -89,10 +89,10 @@ module ActiveRecord
     end
 
     initializer "active_record.database_selector" do
-      if config.active_record.database_selector
+      if options = config.active_record.delete(:database_selector)
         resolver = config.active_record.delete(:database_resolver)
         operations = config.active_record.delete(:database_operations)
-        config.app_middleware.use ActiveRecord::Middleware::DatabaseSelector, resolver, operations
+        config.app_middleware.use ActiveRecord::Middleware::DatabaseSelector, resolver, operations, options
       end
     end
 
diff --git a/activerecord/test/cases/database_selector_test.rb b/activerecord/test/cases/database_selector_test.rb
index 6142e223ce..4106a6ec46 100644
--- a/activerecord/test/cases/database_selector_test.rb
+++ b/activerecord/test/cases/database_selector_test.rb
@@ -95,6 +95,54 @@ module ActiveRecord
       assert @session_store[:last_write]
     end
 
+    def test_read_from_primary_with_options
+      resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session, delay: 5.seconds)
+
+      # 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]
+
+      read = false
+      resolver.read do
+        assert ActiveRecord::Base.connected_to?(role: :writing)
+        read = true
+      end
+      assert read
+    end
+
+    def test_read_from_replica_with_no_delay
+      resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session, delay: 0.seconds)
+
+      # 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]
+
+      read = false
+      resolver.read do
+        assert ActiveRecord::Base.connected_to?(role: :reading)
+        read = true
+      end
+      assert read
+    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)