Merge pull request #51009 from HeyNonster/nony--add-on-all-shards

Add `.shard_keys`, `.sharded?`, & `.connected_to_all_shards` methods to AR Models

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 @@ module ActiveRecord
 
       connections = []
 
+      @shard_keys = shards.keys
+
       if shards.empty?
         shards[:default] = database
       end
@@ -175,6 +177,18 @@ module ActiveRecord
       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 @@ module ActiveRecord
       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