From 77cf5e6d927fd8f5db793be32bb625ac791e5824 Mon Sep 17 00:00:00 2001 From: Nony Dutton Date: Thu, 8 Feb 2024 09:05:22 +0100 Subject: [PATCH] Add `.shard_keys` & `.connected_to_all_shards` Currently, there is no (simple) way to ask a model if it connects to a single database or to multiple shards. Furthermore, without looping through a model's connections, I don't believe there's an easy way to return a list of shards a model can connect to. This commit adds a `@shard_keys` ivar that's set whenever `.connects_to` is called. It sets the ivar to the result of `shards.keys`. `shards` in `.connects_to` defaults to an empty hash and therefore when calling `connects_to database: {...}` `@shard_keys` will be set to an empty array. `@shard_keys` is set _before_ the following lines: ``` if shards.empty? shards[:default] = database end ``` This conditional sets the one and only shard (`:default`) to the value of `database` that we pass to `.connects_to`. This allows for calling `connected_to(shard: :default)` on models configured to only connect to a database e.g.: ```ruby class UnshardedBase < ActiveRecord::Base self.abstract_class = true connects_to database: { writing: :primary } end class UnshardedModel < UnshardedBase end UnshardedBase.connected_to(shard: :default) { UnshardedBase.connection_pool.db_config.name } => primary ``` This is ultimately still an _unsharded_ model which is why `@shard_keys` gets set before the conditional. With the new `@shard_keys` ivar we need a way for descendants of the abstract AR model to return that same value. For that we leverage the existing `.connection_class_for_self` method. That method returns the ancestor of the model where `.connects_to` was called, or returns self if it's the connection class: ```ruby class UnshardedBase < ActiveRecord::Base self.abstract_class = true connects_to database: { writing: :primary } end class UnshardedModel < UnshardedBase end ActiveRecord::Base.connection_class_for_self => ActiveRecord::Base UnshardedBase.connection_class_for_self => UnshardedBase(abstract) UnshardedModel.connection_class_for_self => UnshardedBase(abstract) ``` The new `.shard_keys` method is a getter which returns the value of `@shard_keys` from the connection class or it returns an empty array. The empty array is necessary in cases where `connects_to` was never called. Finally, I've added an `.connected_to_all_shards` method which takes all of the arguments for `.connected_to` except for `shard`. Instead, it loops through every shard key and then delegates everything else to `.connected_to`. I've used `.map` instead of `.each` so that we can collect the results of each block. --- activerecord/CHANGELOG.md | 22 +++ .../lib/active_record/connection_handling.rb | 22 +++ activerecord/test/cases/shard_keys_test.rb | 125 ++++++++++++++++++ 3 files changed, 169 insertions(+) create mode 100644 activerecord/test/cases/shard_keys_test.rb diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md index bb98c7b52d..6232111f82 100644 --- a/activerecord/CHANGELOG.md +++ b/activerecord/CHANGELOG.md @@ -1,3 +1,25 @@ +* Add `.shard_keys`, `.sharded?`, & `.connected_to_all_shards` methods. + + ```ruby + class ShardedBase < ActiveRecord::Base + self.abstract_class = true + + connects_to shards: { + shard_one: { writing: :shard_one }, + shard_two: { writing: :shard_two } + } + end + + class ShardedModel < ShardedBase + end + + ShardedModel.shard_keys => [:shard_one, :shard_two] + ShardedModel.sharded? => true + ShardedBase.connected_to_all_shards { ShardedModel.current_shard } => [:shard_one, :shard_two] + ``` + + *Nony Dutton* + * Optimize `Relation#exists?` when records are loaded and the relation has no conditions. This can avoid queries in some cases. diff --git a/activerecord/lib/active_record/connection_handling.rb b/activerecord/lib/active_record/connection_handling.rb index ff31656c0d..91ddf8285a 100644 --- a/activerecord/lib/active_record/connection_handling.rb +++ b/activerecord/lib/active_record/connection_handling.rb @@ -87,6 +87,8 @@ def connects_to(database: {}, shards: {}) connections = [] + @shard_keys = shards.keys + if shards.empty? shards[:default] = database end @@ -175,6 +177,18 @@ def connected_to_many(*classes, role:, shard: nil, prevent_writes: false) connected_to_stack.pop end + # Passes the block to +connected_to+ for every +shard+ the + # model is configured to connect to (if any), and returns the + # results in an array. + # + # Optionally, +role+ and/or +prevent_writes+ can be passed which + # will be forwarded to each +connected_to+ call. + def connected_to_all_shards(role: nil, prevent_writes: false, &blk) + shard_keys.map do |shard| + connected_to(shard: shard, role: role, prevent_writes: prevent_writes, &blk) + end + end + # Use a specified connection. # # This method is useful for ensuring that a specific connection is @@ -359,6 +373,14 @@ def clear_cache! # :nodoc: connection_pool.schema_cache.clear! end + def shard_keys + connection_class_for_self.instance_variable_get(:@shard_keys) || [] + end + + def sharded? + shard_keys.any? + end + private def resolve_config_for_connection(config_or_env) raise "Anonymous class is not allowed." unless name diff --git a/activerecord/test/cases/shard_keys_test.rb b/activerecord/test/cases/shard_keys_test.rb new file mode 100644 index 0000000000..adde96266e --- /dev/null +++ b/activerecord/test/cases/shard_keys_test.rb @@ -0,0 +1,125 @@ +# frozen_string_literal: true + +require "cases/helper" + +module ActiveRecord + class ShardsKeysTest < ActiveRecord::TestCase + class UnshardedBase < ActiveRecord::Base + self.abstract_class = true + end + + class UnshardedModel < UnshardedBase + end + + class ShardedBase < ActiveRecord::Base + self.abstract_class = true + end + + class ShardedModel < ShardedBase + end + + def setup + ActiveRecord::Base.instance_variable_set(:@shard_keys, nil) + @previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env" + + config = { + "default_env" => { + "primary" => { + adapter: "sqlite3", + database: ":memory:" + }, + "shard_one" => { + adapter: "sqlite3", + database: ":memory:" + }, + "shard_one_reading" => { + adapter: "sqlite3", + database: ":memory:" + }, + "shard_two" => { + adapter: "sqlite3", + database: ":memory:" + }, + "shard_two_reading" => { + adapter: "sqlite3", + database: ":memory:" + }, + } + } + + @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config + + UnshardedBase.connects_to database: { writing: :primary } + + ShardedBase.connects_to shards: { + shard_one: { writing: :shard_one, reading: :shard_one_reading }, + shard_two: { writing: :shard_two, reading: :shard_two_reading }, + } + end + + def teardown + clean_up_connection_handler + ActiveRecord::Base.configurations = @prev_configs + ActiveRecord::Base.establish_connection(:arunit) + ENV["RAILS_ENV"] = @previous_env + end + + def test_connects_to_sets_shard_keys + assert_empty(ActiveRecord::Base.shard_keys) + assert_equal([:shard_one, :shard_two], ShardedBase.shard_keys) + end + + def test_connects_to_sets_shard_keys_for_descendents + assert_equal(ShardedBase.shard_keys, ShardedModel.shard_keys) + end + + def test_sharded? + assert_not ActiveRecord::Base.sharded? + assert_not UnshardedBase.sharded? + assert_not UnshardedModel.sharded? + + assert_predicate ShardedBase, :sharded? + assert_predicate ShardedModel, :sharded? + end + + def test_connected_to_all_shards + unsharded_results = UnshardedBase.connected_to_all_shards do + UnshardedBase.connection_pool.db_config.name + end + + sharded_results = ShardedBase.connected_to_all_shards do + ShardedBase.connection_pool.db_config.name + end + + assert_empty unsharded_results + assert_equal(["shard_one", "shard_two"], sharded_results) + end + + def test_connected_to_all_shards_can_switch_each_to_reading_role + # We teardown the shared connection pool and call .connects_to again + # because .setup_shared_connection_pool overwrites our reading configs + # with the writing role configs. + teardown_shared_connection_pool + ShardedBase.connects_to shards: { + shard_one: { writing: :shard_one, reading: :shard_one_reading }, + shard_two: { writing: :shard_two, reading: :shard_two_reading }, + } + + results = ShardedBase.connected_to_all_shards(role: :reading) do + ShardedBase.connection_pool.db_config.name + end + + assert_equal(["shard_one_reading", "shard_two_reading"], results) + end + + def test_connected_to_all_shards_respects_preventing_writes + assert_not ShardedBase.current_preventing_writes + + results = ShardedBase.connected_to_all_shards(role: :writing, prevent_writes: true) do + ShardedBase.current_preventing_writes + end + + assert_equal([true, true], results) + end + end +end