Add automatic shard swapping middleware

This PR adds a middleware that can be used for automatic shard swapping.
The design is similar to the database selector middleware in that the
resolver is provided by the application to determine which shard to
switch to. The selector also takes options to change the default
behavior of the middleware.

The only supported option is `lock` at the moment which will allow shard
swapping in a request, otherwise it defaults to true and prevents shard
swapping. This will help protect mistakenly switching shards inside of
application code in a multi-tenant application.

The resolver can be designed however the application wants but the basic
idea is that the resolver accesses the `request` headers and uses that
to lookup the subdomain which then looks up the tenant shard name stored
in that table. The tenant table is the "router" for the entire
application and an example resolver looks like this:

```ruby
config.active_record.shard_resolver = ->(request) {
  subdomain = request.subdomain
  tenant = Tenant.find_by_subdomain!(subdomain)
  tenant.shard
}
```

The `Tenant` table in this example would have `subdomain` and `shard`
attributes that are inserted into the database. These are used to route
to the shard in `connected_to`. Ie if we had a `Tenant` with the
subdomain `github` and the shard name `github_shard` we'd lookup the
connection with `ActiveRecord::Base.connected_to(shard: :github_shard)
{}` and all queries within that block (request) would be scoped to
the github shard.

Co-authored-by: John Crepezzi <[email protected]>

activerecord/CHANGELOG.md

@@ -1,3 +1,19 @@
+*   Add middleware for automatic shard swapping.
+
+    Provides a basic middleware to perform automatic shard swapping. Applications will provide a resolver which will determine for an individual request which shard should be used. Example:
+
+    ```ruby
+    config.active_record.shard_resolver = ->(request) {
+      subdomain = request.subdomain
+      tenant = Tenant.find_by_subdomain!(subdomain)
+      tenant.shard
+    }
+    ```
+
+    See guides for more details.
+
+    *Eileen M. Uchitelle*, *John Crepezzi*
+
 *   Remove deprecated support to pass a column to `type_cast`.
 
     *Rafael Mendonça França*

activerecord/lib/active_record.rb

@@ -155,6 +155,7 @@ module Middleware
     extend ActiveSupport::Autoload
 
     autoload :DatabaseSelector, "active_record/middleware/database_selector"
+    autoload :ShardSelector, "active_record/middleware/shard_selector"
   end
 
   module Tasks

activerecord/lib/active_record/connection_handling.rb

@@ -211,11 +211,14 @@ def connecting_to(role: default_role, shard: default_shard, prevent_writes: fals
     # nested call to connected_to or connected_to_many to swap again. This
     # is useful in cases you're using sharding to provide per-request
     # database isolation.
-    def prohibit_shard_swapping
-      Thread.current.thread_variable_set(:prohibit_shard_swapping, true)
+    def prohibit_shard_swapping(enabled = true)
+      prev_value = Thread.current.thread_variable_get(:prohibit_shard_swapping)
+
+      Thread.current.thread_variable_set(:prohibit_shard_swapping, enabled)
+
       yield
     ensure
-      Thread.current.thread_variable_set(:prohibit_shard_swapping, false)
+      Thread.current.thread_variable_set(:prohibit_shard_swapping, prev_value)
     end
 
     # Determine whether or not shard swapping is currently prohibited

activerecord/lib/active_record/core.rb

@@ -77,6 +77,8 @@ def self.configurations
 
       class_attribute :default_shard, instance_writer: false
 
+      class_attribute :shard_selector, instance_accessor: false, default: nil
+
       def self.application_record_class? # :nodoc:
         if ActiveRecord.application_record_class
           self == ActiveRecord.application_record_class

activerecord/lib/active_record/middleware/shard_selector.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+require "active_record/middleware/database_selector/resolver"
+
+module ActiveRecord
+  module 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
+    #   }
+    class ShardSelector
+      def initialize(app, resolver, options = {})
+        @app = app
+        @resolver = resolver
+        @options = options
+      end
+
+      attr_reader :resolver, :options
+
+      def call(env)
+        request = ActionDispatch::Request.new(env)
+
+        shard = selected_shard(request)
+
+        set_shard(shard) do
+          @app.call(env)
+        end
+      end
+
+      private
+        def selected_shard(request)
+          resolver.call(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)
+          end
+        end
+    end
+  end
+end

activerecord/lib/active_record/railtie.rb

@@ -102,6 +102,14 @@ class Railtie < Rails::Railtie # :nodoc:
       end
     end
 
+    initializer "active_record.shard_selector" do
+      if resolver = config.active_record.shard_resolver
+        options = config.active_record.shard_selector || {}
+
+        config.app_middleware.use ActiveRecord::Middleware::ShardSelector, resolver, options
+      end
+    end
+
     initializer "Check for cache versioning support" do
       config.after_initialize do |app|
         ActiveSupport.on_load(:active_record) do
@@ -232,6 +240,8 @@ class Railtie < Rails::Railtie # :nodoc:
           :database_selector,
           :database_resolver,
           :database_resolver_context,
+          :shard_selector,
+          :shard_resolver,
           :query_log_tags_enabled,
           :query_log_tags,
           :cache_query_log_tags,

activerecord/test/cases/shard_selector_test.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require "cases/helper"
+require "models/person"
+require "action_dispatch"
+
+module ActiveRecord
+  class ShardSelectorTest < ActiveRecord::TestCase
+    def test_middleware_locks_to_shard_by_default
+      middleware = ActiveRecord::Middleware::ShardSelector.new(lambda { |env|
+        assert_predicate ActiveRecord::Base, :shard_swapping_prohibited?
+        [200, {}, ["body"]]
+      }, ->(*) { :shard_one })
+
+      assert_equal [200, {}, ["body"]], middleware.call("REQUEST_METHOD" => "GET")
+    end
+
+    def test_middleware_can_turn_off_lock_option
+      middleware = ActiveRecord::Middleware::ShardSelector.new(lambda { |env|
+        assert_not_predicate ActiveRecord::Base, :shard_swapping_prohibited?
+        [200, {}, ["body"]]
+      }, ->(*) { :shard_one }, { lock: false })
+
+      assert_equal [200, {}, ["body"]], middleware.call("REQUEST_METHOD" => "GET")
+    end
+
+    def test_middleware_can_change_shards
+      middleware = ActiveRecord::Middleware::ShardSelector.new(lambda { |env|
+        assert ActiveRecord::Base.connected_to?(role: :writing, shard: :shard_one)
+        [200, {}, ["body"]]
+      }, ->(*) { :shard_one })
+
+      assert_equal [200, {}, ["body"]], middleware.call("REQUEST_METHOD" => "GET")
+    end
+
+    def test_middleware_can_handle_string_shards
+      middleware = ActiveRecord::Middleware::ShardSelector.new(lambda { |env|
+        assert ActiveRecord::Base.connected_to?(role: :writing, shard: :shard_one)
+        [200, {}, ["body"]]
+      }, ->(*) { "shard_one" })
+
+      assert_equal [200, {}, ["body"]], middleware.call("REQUEST_METHOD" => "GET")
+    end
+  end
+end

guides/source/active_record_multiple_databases.md

@@ -31,7 +31,6 @@ databases
 
 The following features are not (yet) supported:
 
-* Automatic swapping for horizontal sharding
 * Load balancing replicas
 
 ## Setting up your application
@@ -276,7 +275,7 @@ $ bin/rails generate scaffold Dog name:string --database animals --parent Animal
 This will skip generating `AnimalsRecord` since you've indicated to Rails that you want to
 use a different parent class.
 
-## Activating automatic connection switching
+## Activating automatic role switching
 
 Finally, in order to use the read-only replica in your application, you'll need to activate
 the middleware for automatic switching.
@@ -426,6 +425,40 @@ ActiveRecord::Base.connected_to(role: :reading, shard: :shard_one) do
 end
 ```
 
+## Activating automatic shard switching
+
+Applications are able to automatically switch shards per request using the provided
+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:
+
+```ruby
+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:
+
+```ruby
+config.active_record.shard_resolver = ->(request) {
+  subdomain = request.subdomain
+  tenant = Tenant.find_by_subdomain!(subdomain)
+  tenant.shard
+}
+```
+
 ## Migrate to the new connection handling
 
 In Rails 6.1+, Active Record provides a new internal API for connection management.

railties/lib/rails/generators/rails/app/templates/config/environments/production.rb.tt

@@ -125,4 +125,14 @@ Rails.application.configure do
   # config.active_record.database_selector = { delay: 2.seconds }
   # config.active_record.database_resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver
   # config.active_record.database_resolver_context = ActiveRecord::Middleware::DatabaseSelector::Resolver::Session
+
+  # Inserts middleware to perform automatic shard swapping. The `shard_selector` hash
+  # can be used to pass options to the `ShardSelector` middleware. The `lock` option is
+  # used to determine whether shard swapping should be prohibited for the request.
+  #
+  # The `shard_resolver` option is used by the middleware to determine which shard
+  # to switch to. The application must provide a mechanism for finding the shard name
+  # in a proc. See guides for an example.
+  # config.active_record.shard_selector = { lock: true }
+  # config.active_record.shard_resolver = ->(request) { Tenant.find_by!(host: request.host).shard }
 end

railties/test/application/middleware_test.rb

@@ -345,6 +345,14 @@ def index
       assert_equal "/foo/?something", env["ORIGINAL_FULLPATH"]
     end
 
+    test "shard selector middleware is installed by config option" do
+      add_to_config "config.active_record.shard_resolver = ->(*) { }"
+
+      boot!
+
+      assert_includes middleware, "ActiveRecord::Middleware::ShardSelector"
+    end
+
     private
       def boot!
         require "#{app_path}/config/environment"