From 0bc8c03034aade0eb3e6bb754249d6f38007032b Mon Sep 17 00:00:00 2001 From: Mike Dalessio Date: Fri, 29 Nov 2024 12:42:10 -0500 Subject: [PATCH] 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 | 12 ++++ .../middleware/shard_selector.rb | 68 ++++++++++++------- .../test/cases/shard_selector_test.rb | 33 +++++++++ .../active_record_multiple_databases.md | 44 ++++++------ 4 files changed, 112 insertions(+), 45 deletions(-) diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md index f6aa762a6fb9c..b3aa936718284 100644 --- a/activerecord/CHANGELOG.md +++ b/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. diff --git a/activerecord/lib/active_record/middleware/shard_selector.rb b/activerecord/lib/active_record/middleware/shard_selector.rb index 1db88b904fb44..5764ab85a4556 100644 --- a/activerecord/lib/active_record/middleware/shard_selector.rb +++ b/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 diff --git a/activerecord/test/cases/shard_selector_test.rb b/activerecord/test/cases/shard_selector_test.rb index 0d808fb309fdf..1074d83f918d7 100644 --- a/activerecord/test/cases/shard_selector_test.rb +++ b/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 diff --git a/guides/source/active_record_multiple_databases.md b/guides/source/active_record_multiple_databases.md index 5deca2cd07cf5..45e78158c4e4c 100644 --- a/guides/source/active_record_multiple_databases.md +++ b/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