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.

Author: HeyNonster

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.

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

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