From 27d51da866c5e8d59eef159cf78f63b022e6314e Mon Sep 17 00:00:00 2001 From: eileencodes Date: Mon, 2 Nov 2020 08:26:46 -0500 Subject: [PATCH] Add `connected_to_many` API Now that we have implemented granular connection swapping in #40370 we need a new API that will allow connections to multiple databases. The reason we need this API is it will prevent deep nesting in cases where we know that we want 3 of our 5 databases to connect to reading and leave the rest on writing. With this API, instead of writing: ```ruby AnimalsRecord.connected_to(role: :reading) do MealsRecord.connected_to(role: :reading) do Dog.first # read from animals replica Dinner.first # read from meals replica Person.first # read from primary writer end end ``` This API would allow you to write: ```ruby ActiveRecord::Base.connected_to_many([AnimalsRecord, MealsRecord], role: :reading) do Dog.first # read from animals replica Dinner.first # read from meals replica Person.first # read from primary writer end ``` This would come in especially handy for deeper nesting past 2 databases. Co-authored-by: John Crepezzi --- activerecord/CHANGELOG.md | 24 ++++++ .../lib/active_record/connection_handling.rb | 36 +++++++- activerecord/lib/active_record/core.rb | 13 ++- activerecord/test/cases/base_test.rb | 30 +++++++ .../connection_swapping_nested_test.rb | 83 +++++++++++++++++++ 5 files changed, 176 insertions(+), 10 deletions(-) diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md index 41dd28987d..ea8581cc1a 100644 --- a/activerecord/CHANGELOG.md +++ b/activerecord/CHANGELOG.md @@ -1,3 +1,27 @@ +* Add `connected_to_many` API. + + This API allows applications to connect to multiple databases at once without switching all of them or implementing a deeply nested stack. + + Before: + + AnimalsRecord.connected_to(role: :reading) do + MealsRecord.connected_to(role: :reading) do + Dog.first # read from animals replica + Dinner.first # read from meals replica + Person.first # read from primary writer + end + end + + After: + + ActiveRecord::Base.connected_to_many([AnimalsRecord, MealsRecord], role: :reading) do + Dog.first # read from animals replica + Dinner.first # read from meals replica + Person.first # read from primary writer + end + + *Eileen M. Uchitelle*, *John Crepezzi* + * Add option to raise or log for `ActiveRecord::StrictLoadingViolationError`. Some applications may not want to raise an error in production if using `strict_loading`. This would allow an application to set strict loading to log for the production environment while still raising in development and test environments. diff --git a/activerecord/lib/active_record/connection_handling.rb b/activerecord/lib/active_record/connection_handling.rb index f4e04c9c54..83c887defe 100644 --- a/activerecord/lib/active_record/connection_handling.rb +++ b/activerecord/lib/active_record/connection_handling.rb @@ -172,6 +172,36 @@ module ActiveRecord end end + # Connects a role and/or shard to the provided connection names. Optionally `prevent_writes` + # can be passed to block writes on a connection. `reading` will automatically set + # `prevent_writes` to true. + # + # `connected_to_many` is an alternative to deeply nested `connected_to` blocks. + # + # Usage: + # + # ActiveRecord::Base.connected_to(AnimalsRecord, MealsRecord], role: :reading) do + # Dog.first # Read from animals replica + # Dinner.first # Read from meals replica + # Person.first # Read from primary writer + # end + def connected_to_many(classes, role:, shard: nil, prevent_writes: false) + if legacy_connection_handling + raise NotImplementedError, "connected_to_many is not available with legacy connection handling" + end + + if self != Base || classes.include?(Base) + raise NotImplementedError, "connected_to_many can only be called on ActiveRecord::Base." + end + + prevent_writes = true if role == reading_role + + connected_to_stack << { role: role, shard: shard, prevent_writes: prevent_writes, klasses: classes } + yield + ensure + connected_to_stack.pop + end + # Use a specified connection. # # This method is useful for ensuring that a specific connection is @@ -186,7 +216,7 @@ module ActiveRecord prevent_writes = true if role == reading_role - self.connected_to_stack << { role: role, shard: shard, prevent_writes: prevent_writes, klass: self } + self.connected_to_stack << { role: role, shard: shard, prevent_writes: prevent_writes, klasses: [self] } end # Prevent writing to the database regardless of role. @@ -341,12 +371,12 @@ module ActiveRecord if ActiveRecord::Base.legacy_connection_handling with_handler(role.to_sym) do connection_handler.while_preventing_writes(prevent_writes) do - self.connected_to_stack << { shard: shard, klass: self } + self.connected_to_stack << { shard: shard, klasses: [self] } yield end end else - self.connected_to_stack << { role: role, shard: shard, prevent_writes: prevent_writes, klass: self } + self.connected_to_stack << { role: role, shard: shard, prevent_writes: prevent_writes, klasses: [self] } return_value = yield return_value.load if return_value.is_a? ActiveRecord::Relation return_value diff --git a/activerecord/lib/active_record/core.rb b/activerecord/lib/active_record/core.rb index 543721da1e..c1fdfbc2c3 100644 --- a/activerecord/lib/active_record/core.rb +++ b/activerecord/lib/active_record/core.rb @@ -187,15 +187,14 @@ module ActiveRecord connection_handlers.key(connection_handler) || default_role else connected_to_stack.reverse_each do |hash| - return hash[:role] if hash[:role] && hash[:klass] == Base - return hash[:role] if hash[:role] && hash[:klass] == abstract_base_class + return hash[:role] if hash[:role] && hash[:klasses].include?(Base) + return hash[:role] if hash[:role] && hash[:klasses].include?(abstract_base_class) end default_role end end - # Returns the symbol representing the current connected shard. # # ActiveRecord::Base.connected_to(role: :reading) do @@ -207,8 +206,8 @@ module ActiveRecord # end def self.current_shard connected_to_stack.reverse_each do |hash| - return hash[:shard] if hash[:shard] && hash[:klass] == Base - return hash[:shard] if hash[:shard] && hash[:klass] == abstract_base_class + return hash[:shard] if hash[:shard] && hash[:klasses].include?(Base) + return hash[:shard] if hash[:shard] && hash[:klasses].include?(abstract_base_class) end default_shard @@ -229,8 +228,8 @@ module ActiveRecord connection_handler.prevent_writes else connected_to_stack.reverse_each do |hash| - return hash[:prevent_writes] if !hash[:prevent_writes].nil? && hash[:klass] == Base - return hash[:prevent_writes] if !hash[:prevent_writes].nil? && hash[:klass] == abstract_base_class + return hash[:prevent_writes] if !hash[:prevent_writes].nil? && hash[:klasses].include?(Base) + return hash[:prevent_writes] if !hash[:prevent_writes].nil? && hash[:klasses].include?(abstract_base_class) end false diff --git a/activerecord/test/cases/base_test.rb b/activerecord/test/cases/base_test.rb index 9f0f65498a..e639ad6e45 100644 --- a/activerecord/test/cases/base_test.rb +++ b/activerecord/test/cases/base_test.rb @@ -1688,4 +1688,34 @@ class BasicsTest < ActiveRecord::TestCase ensure ActiveRecord::Base.legacy_connection_handling = old_value end + + test "#connected_to_many doesn't work with legacy connection handling" do + old_value = ActiveRecord::Base.legacy_connection_handling + ActiveRecord::Base.legacy_connection_handling = true + + assert_raises NotImplementedError do + ActiveRecord::Base.connected_to_many([AbstractCompany], role: :writing) + end + ensure + ActiveRecord::Base.legacy_connection_handling = old_value + end + + test "#connected_to_many cannot be called on anything but ActiveRecord::Base" do + assert_raises NotImplementedError do + AbstractCompany.connected_to_many([AbstractCompany], role: :writing) + end + end + + test "#connected_to_many cannot be called with classes that include ActiveRecord::Base" do + assert_raises NotImplementedError do + ActiveRecord::Base.connected_to_many([ActiveRecord::Base], role: :writing) + end + end + + test "#connected_to_many sets prevent_writes if role is reading" do + ActiveRecord::Base.connected_to_many([AbstractCompany], role: :reading) do + assert AbstractCompany.current_preventing_writes + assert_not ActiveRecord::Base.current_preventing_writes + end + end end diff --git a/activerecord/test/cases/connection_adapters/connection_swapping_nested_test.rb b/activerecord/test/cases/connection_adapters/connection_swapping_nested_test.rb index 844eba6509..d53026adb5 100644 --- a/activerecord/test/cases/connection_adapters/connection_swapping_nested_test.rb +++ b/activerecord/test/cases/connection_adapters/connection_swapping_nested_test.rb @@ -28,6 +28,13 @@ module ActiveRecord class SecondaryModel < SecondaryBase end + class TertiaryBase < ActiveRecord::Base + self.abstract_class = true + end + + class TertiaryModel < TertiaryBase + end + unless in_memory_db? def test_roles_can_be_swapped_granularly previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env" @@ -266,6 +273,82 @@ module ActiveRecord ActiveRecord::Base.establish_connection(:arunit) ENV["RAILS_ENV"] = previous_env end + + def test_connected_to_many + previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env" + + config = { + "default_env" => { + "primary" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3" }, + "primary_replica" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3", "replica" => true }, + "primary_shard_one" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_one.sqlite3" }, + "primary_shard_one_replica" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_one.sqlite3", "replica" => true }, + "primary_shard_two" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_two.sqlite3" }, + "primary_shard_two_replica" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_two.sqlite3", "replica" => true }, + "secondary" => { "adapter" => "sqlite3", "database" => "test/db/secondary.sqlite3" }, + "secondary_replica" => { "adapter" => "sqlite3", "database" => "test/db/secondary.sqlite3", "replica" => true }, + "secondary_shard_one" => { "adapter" => "sqlite3", "database" => "test/db/secondary_shard_one.sqlite3" }, + "secondary_shard_one_replica" => { "adapter" => "sqlite3", "database" => "test/db/secondary_shard_one.sqlite3", "replica" => true }, + "secondary_shard_two" => { "adapter" => "sqlite3", "database" => "test/db/secondary_shard_two.sqlite3" }, + "secondary_shard_two_replica" => { "adapter" => "sqlite3", "database" => "test/db/secondary_shard_two.sqlite3", "replica" => true }, + "tertiary" => { "adapter" => "sqlite3", "database" => "test/db/tertiary.sqlite3" }, + "tertiary_replica" => { "adapter" => "sqlite3", "database" => "test/db/tertiary.sqlite3", "replica" => true }, + "tertiary_shard_one" => { "adapter" => "sqlite3", "database" => "test/db/tertiary_shard_one.sqlite3" }, + "tertiary_shard_one_replica" => { "adapter" => "sqlite3", "database" => "test/db/tertiary_shard_one.sqlite3", "replica" => true }, + "tertiary_shard_two" => { "adapter" => "sqlite3", "database" => "test/db/tertiary_shard_two.sqlite3" }, + "tertiary_shard_two_replica" => { "adapter" => "sqlite3", "database" => "test/db/tertiary_shard_two.sqlite3", "replica" => true } + } + } + + @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config + + PrimaryBase.connects_to(shards: { + default: { writing: :primary, reading: :primary_replica }, + shard_one: { writing: :primary_shard_one, reading: :primary_shard_one_replica } + }) + + SecondaryBase.connects_to(shards: { + default: { writing: :secondary, reading: :secondary_replica }, + shard_one: { writing: :secondary_shard_one, reading: :secondary_shard_one_replica }, + shard_two: { writing: :secondary_shard_two, reading: :secondary_shard_two_replica } + }) + + TertiaryBase.connects_to(shards: { + default: { writing: :tertiary, reading: :tertiary_replica }, + shard_one: { writing: :tertiary_shard_one, reading: :tertiary_shard_one_replica }, + shard_two: { writing: :tertiary_shard_two, reading: :tertiary_shard_two_replica } + }) + + # Switch everything to writing, default shard + ActiveRecord::Base.connected_to(role: :writing, shard: :default) do + assert_equal "primary", PrimaryBase.connection_pool.db_config.name + assert_equal "secondary", SecondaryBase.connection_pool.db_config.name + assert_equal "tertiary", TertiaryBase.connection_pool.db_config.name + + # Switch two to reading + ActiveRecord::Base.connected_to_many([SecondaryBase, TertiaryBase], role: :reading) do + assert_equal "primary", PrimaryBase.connection_pool.db_config.name + assert_equal "secondary_replica", SecondaryBase.connection_pool.db_config.name + assert_equal "tertiary_replica", TertiaryBase.connection_pool.db_config.name + + # Switch one back + ActiveRecord::Base.connected_to_many([TertiaryBase], role: :writing) do + assert_equal "primary", PrimaryBase.connection_pool.db_config.name + assert_equal "secondary_replica", SecondaryBase.connection_pool.db_config.name + assert_equal "tertiary", TertiaryBase.connection_pool.db_config.name + end + end + + # Switched back + assert_equal "primary", PrimaryBase.connection_pool.db_config.name + assert_equal "secondary", SecondaryBase.connection_pool.db_config.name + assert_equal "tertiary", TertiaryBase.connection_pool.db_config.name + end + ensure + ActiveRecord::Base.configurations = @prev_configs + ActiveRecord::Base.establish_connection(:arunit) + ENV["RAILS_ENV"] = previous_env + end end end end