ShardSelector supports connecting to a non-primary database.

This allows for the ShardSelector to be used for a non-primary sharded
database. Previously only the primary database (i.e.,
ActiveRecord::Base) was supported.

Also updates the Multiple Databases guide and improves the
ShardSelector documentation.

activerecord/CHANGELOG.md

@@ -69,4 +69,16 @@
 
     *Kazuma Watanabe*
 
+*   `ActiveRecord::Middleware::ShardSelector` supports automatic switching on non-primary databases.
+
+    A new configuration option, `class_name:`, is introduced to
+    `config.active_record.shard_selector` to allow a non-primary sharded database to use the shard
+    selection middleware:
+
+    ```
+    config.active_record.shard_selector = { class_name: "AnimalsRecord" }
+    ```
+
+    *Mike Dalessio*
+
 Please check [8-0-stable](https://github.com/rails/rails/blob/8-0-stable/activerecord/CHANGELOG.md) for previous changes.

activerecord/lib/active_record/middleware/shard_selector.rb

@@ -4,30 +4,44 @@ module ActiveRecord
   module Middleware
     # = Shard Selector \Middleware
     #
-    # The ShardSelector Middleware provides a framework for automatically
-    # swapping shards. \Rails provides a basic framework to determine which
-    # shard to switch to and allows for applications to write custom strategies
-    # for swapping if needed.
-    #
-    # The ShardSelector takes a set of options (currently only +lock+ is supported)
-    # that can be used by the middleware to alter behavior. +lock+ is
-    # true by default and will prohibit the request from switching shards once
-    # inside the block. If +lock+ is false, then shard swapping will be allowed.
-    # For tenant based sharding, +lock+ should always be true to prevent application
-    # code from mistakenly switching between tenants.
-    #
-    # Options can be set in the config:
-    #
-    #   config.active_record.shard_selector = { lock: true }
-    #
-    # Applications must also provide the code for the resolver as it depends on application
-    # specific models. An example resolver would look like this:
-    #
-    #   config.active_record.shard_resolver = ->(request) {
-    #     subdomain = request.subdomain
-    #     tenant = Tenant.find_by_subdomain!(subdomain)
-    #     tenant.shard
-    #   }
+    # The ShardSelector Middleware provides a framework for automatically swapping shards. \Rails
+    # provides a basic framework to determine which shard to switch to and allows for applications
+    # to write custom strategies for swapping if needed.
+    #
+    # == Setup
+    #
+    # Applications must provide a resolver that will provide application-specific logic for
+    # selecting the appropriate shard. Setting +config.active_record.shard_resolver+ will cause
+    # Rails to add ShardSelector to the default middleware stack.
+    #
+    # The resolver, along with any configuration options, can be set in the application
+    # configuration using an initializer like so:
+    #
+    #   Rails.application.configure do
+    #     config.active_record.shard_selector = { lock: false, class_name: "AnimalsRecord" }
+    #     config.active_record.shard_resolver = ->(request) {
+    #       subdomain = request.subdomain
+    #       tenant = Tenant.find_by_subdomain!(subdomain)
+    #       tenant.shard
+    #     }
+    #   end
+    #
+    # == Configuration
+    #
+    # The behavior of ShardSelector can be altered through some configuration options.
+    #
+    # [+lock:+]
+    #   +lock+ is true by default and will prohibit the request from switching shards once inside the
+    #   block. If +lock+ is false, then shard swapping will be allowed. For tenant based sharding,
+    #   +lock+ should always be true to prevent application code from mistakenly switching between
+    #   tenants.
+    #
+    # [+class_name:+]
+    #   +class_name+ is the name of the abstract base class used for sharding. By default, the
+    #   ShardSelector will use ActiveRecord::Base, but if the application is sharding a non-primary
+    #   database, then this option should be set to the name of that database's abstract base
+    #   class.
+    #
     class ShardSelector
       def initialize(app, resolver, options = {})
         @app = app
@@ -53,8 +67,10 @@ def selected_shard(request)
         end
 
         def set_shard(shard, &block)
-          ActiveRecord::Base.connected_to(shard: shard.to_sym) do
-            ActiveRecord::Base.prohibit_shard_swapping(options.fetch(:lock, true), &block)
+          klass = options[:class_name]&.constantize || ActiveRecord::Base
+
+          klass.connected_to(shard: shard.to_sym) do
+            klass.prohibit_shard_swapping(options.fetch(:lock, true), &block)
           end
         end
     end

activerecord/test/cases/shard_selector_test.rb

@@ -41,5 +41,38 @@ def test_middleware_can_handle_string_shards
 
       assert_equal [200, {}, ["body"]], middleware.call("REQUEST_METHOD" => "GET")
     end
+
+    def test_middleware_can_connect_to_non_primary_database_model
+      klass = Class.new(ActiveRecord::Base) do |k|
+        class << self
+          attr_reader :connected_to_shard
+
+          def connected_to(shard:)
+            @connected_to_shard = shard
+            yield
+          end
+
+          def prohibit_shard_swapping(...)
+            yield
+          end
+
+          def connected_to?(role: nil, shard:)
+            @connected_to_shard.to_sym == shard.to_sym
+          end
+        end
+      end
+      Object.const_set :ShardSelectorTestModel, klass
+
+      middleware = ActiveRecord::Middleware::ShardSelector.new(lambda { |env|
+        assert_not ActiveRecord::Base.connected_to?(role: :writing, shard: :shard_one)
+        assert klass.connected_to?(role: :writing, shard: :shard_one)
+        [200, {}, ["body"]]
+      }, ->(*) { :shard_one }, { class_name: "ShardSelectorTestModel" })
+
+      assert_equal [200, {}, ["body"]], middleware.call("REQUEST_METHOD" => "GET")
+      assert_equal(:shard_one, klass.connected_to_shard)
+    ensure
+      Object.send(:remove_const, :ShardSelectorTestModel)
+    end
   end
 end

guides/source/active_record_multiple_databases.md

@@ -493,29 +493,18 @@ end
 
 ## Activating Automatic Shard Switching
 
-Applications are able to automatically switch shards per request using the provided
-middleware.
+Applications are able to automatically switch shards per request using the `ShardSelector`
+middleware, which allows an application to provide custom logic for determining the appropriate
+shard for each request.
 
-The `ShardSelector` middleware provides a framework for automatically
-swapping shards. Rails provides a basic framework to determine which
-shard to switch to and allows for applications to write custom strategies
-for swapping if needed.
-
-`ShardSelector` takes a set of options (currently only `lock` is supported)
-that can be used by the middleware to alter behavior. `lock` is
-true by default and will prohibit the request from switching shards once
-inside the block. If `lock` is false, then shard swapping will be allowed.
-For tenant based sharding, `lock` should always be true to prevent application
-code from mistakenly switching between tenants.
-
-The same generator as the database selector can be used to generate the file for
-automatic shard swapping:
+The same generator used for the database selector above can be used to generate an initializer file
+for automatic shard swapping:
 
 ```bash
 $ bin/rails g active_record:multi_db
 ```
 
-Then in the generated `config/initializers/multi_db.rb` uncomment the following:
+Then in the generated `config/initializers/multi_db.rb` uncomment and modify the following code:
 
 ```ruby
 Rails.application.configure do
@@ -524,8 +513,8 @@ Rails.application.configure do
 end
 ```
 
-Applications must provide the code for the resolver as it depends on application
-specific models. An example resolver would look like this:
+Applications must provide a resolver to provide application-specific logic. An example resolver that
+uses subdomain to determine the shard might look like this:
 
 ```ruby
 config.active_record.shard_resolver = ->(request) {
@@ -535,6 +524,23 @@ config.active_record.shard_resolver = ->(request) {
 }
 ```
 
+The behavior of `ShardSelector` can be altered through some configuration options.
+
+`lock` is true by default and will prohibit the request from switching shards during the request. If
+`lock` is false, then shard swapping will be allowed. For tenant-based sharding, `lock` should
+always be true to prevent application code from mistakenly switching between tenants.
+
+`class_name` is the name of the abstract base class used for sharding. By default, the
+`ShardSelector` will use `ActiveRecord::Base`, but if the application is sharding a non-primary
+database, then this option should be set to the name of that database's abstract base class.
+
+Options may be set in the application configuration:
+
+``` ruby
+config.active_record.shard_selector = { lock: true, class_name: "AnimalsRecord" }
+```
+
+
 ## Granular Database Connection Switching
 
 Starting from Rails 6.1, it's possible to switch connections for one database