ShardSelector supports granular database connection switching

This allows for the ShardSelector to be used for granular switching on
a specific abstract connection class. Previously switching was done
globally on ActiveRecord::Base.

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

activerecord/CHANGELOG.md

@@ -1,3 +1,19 @@
+*   `ActiveRecord::Middleware::ShardSelector` supports granular database connection switching.
+
+    A new configuration option, `class_name:`, is introduced to
+    `config.active_record.shard_selector` to allow an application to specify the abstract connection
+    class to be switched by the shard selection middleware. The default class is
+    `ActiveRecord::Base`.
+
+    For example, this configuration tells `ShardSelector` to switch shards using
+    `AnimalsRecord.connected_to`:
+
+    ```
+    config.active_record.shard_selector = { class_name: "AnimalsRecord" }
+    ```
+
+    *Mike Dalessio*
+
 *   Reset relations after `insert_all`/`upsert_all`.
 
     Bulk insert/upsert methods will now call `reset` if used on a relation, matching the behavior of `update_all`.

activerecord/lib/active_record/middleware/shard_selector.rb

@@ -9,25 +9,40 @@ module Middleware
     # 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.
+    # == Setup
     #
-    # Options can be set in the config:
+    # 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.
     #
-    #   config.active_record.shard_selector = { lock: true }
+    # The resolver, along with any configuration options, can be set in the application
+    # configuration using an initializer like so:
     #
-    # Applications must also provide the code for the resolver as it depends on application
-    # specific models. An example resolver would look like this:
+    #   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 switching 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 connection class to switch. By
+    #   default, the ShardSelector will use ActiveRecord::Base, but if the
+    #   application has multiple databases, then this option should be set to
+    #   the name of the sharded database's abstract connection class.
     #
-    #   config.active_record.shard_resolver = ->(request) {
-    #     subdomain = request.subdomain
-    #     tenant = Tenant.find_by_subdomain!(subdomain)
-    #     tenant.shard
-    #   }
     class ShardSelector
       def initialize(app, resolver, options = {})
         @app = app
@@ -53,8 +68,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_do_granular_database_connection_switching
+      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,25 @@ 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 connection class to switch. By default, the `ShardSelector`
+will use `ActiveRecord::Base`, but if the application has multiple databases, then this option
+should be set to the name of the sharded database's abstract connection class.
+
+Options may be set in the application configuration. For example, this configuration tells
+`ShardSelector` to switch shards using `AnimalsRecord.connected_to`:
+
+
+``` 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