Merge pull request rails#43665 from eileencodes/add-shard-middleware

Add automatic shard swapping middleware

Author: eileencodes

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"