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 <john.crepezzi@gmail.com>

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.

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

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

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

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