Skip to content

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

Rodrigo Muino Tomonari requested to merge github/fork/HeyNonster/nony--add-on-all-shards into main

Created by: HeyNonster

Motivation / Background

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.

Detail

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.:

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:

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 .on_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.

Checklist

Before submitting the PR make sure the following are checked:

  • This Pull Request is related to one change. Changes that are unrelated should be opened in separate PRs.
  • Commit message has a detailed description of what changed and why. If this PR fixes a related issue include it in the commit message. Ex: [Fix #issue-number]
  • Tests are added or updated if you fix a bug or add a feature.
  • CHANGELOG files are updated for the changed libraries if there is a behavior change or additional feature. Minor bug fixes and documentation changes should not be included.

Merge request reports