Skip to content

Commit 106d5c5

Browse files
HeyNonsterbyroot
authored andcommitted
Support integer shard keys
Currently, `.connects_to(shards: ...)` coerces all of the shard keys into symbols. Because an integer cannot be coerced into a symbol, attempting to identify a shard with an integer will raise an exception: ```ruby 1.to_sym (irb):1:in '<main>': undefined method 'to_sym' for an instance of Integer (NoMethodError) ``` I think it would be useful to support integers as shard keys along with symbols. This would simplify shard switching when shards are identified by integer columns in a database. As an example, if there is an `accounts` table with a `shard` integer column which identifies which shard that account's data lives on, this currently would not work: ```ruby account = Account.first ActiveRecord::Base.connected_to(shard: account.shard) do # Raises NoMethodError end ``` A workaround would be to first coerce or serialize the integer into a string but that's unideal: ```ruby account = Account.first ActiveRecord::Base.connected_to(shard: account.shard.to_s.to_sym) do # ... end ``` This makes a `.connects_to` change, coercing the shard key into a symbol _unless_ the key is an integer. From there, passing a symbol or integer to `connected_to(shard: ...)` should both work. In the test, we call: ```ruby ActiveRecord::Base.default_shard = 0 ``` Normally, it would be `:default`. In that case we'd have to mix symbols and integers in the `.connects_to:(shards:)` hash which is pretty ugly.
1 parent cea7cae commit 106d5c5

File tree

3 files changed

+78
-1
lines changed

3 files changed

+78
-1
lines changed

activerecord/CHANGELOG.md

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,18 @@
1+
* Add support for integer shard keys.
2+
```ruby
3+
# Now accepts symbols as shard keys.
4+
ActiveRecord::Base.connects_to(shards: {
5+
1: { writing: :primary_shard_one, reading: :primary_shard_one },
6+
2: { writing: :primary_shard_two, reading: :primary_shard_two},
7+
})
8+
9+
ActiveRecord::Base.connected_to(shard: 1) do
10+
# ..
11+
end
12+
```
13+
14+
*Nony Dutton*
15+
116
* Add `ActiveRecord::Base.only_columns`
217

318
Similar in use case to `ignored_columns` but listing columns to consider rather than the ones

activerecord/lib/active_record/connection_handling.rb

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,8 @@ def connects_to(database: {}, shards: {})
100100
db_config = resolve_config_for_connection(database_key)
101101

102102
self.connection_class = true
103-
connections << connection_handler.establish_connection(db_config, owner_name: self, role: role, shard: shard.to_sym)
103+
shard = shard.to_sym unless shard.is_a? Integer
104+
connections << connection_handler.establish_connection(db_config, owner_name: self, role: role, shard: shard)
104105
end
105106
end
106107

activerecord/test/cases/connection_adapters/connection_handlers_sharding_db_test.rb

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,67 @@ def test_establish_connection_using_3_levels_config_with_shards_and_replica
103103
ENV["RAILS_ENV"] = previous_env
104104
end
105105

106+
class IntegerKeysBase < ActiveRecord::Base
107+
self.abstract_class = true
108+
end
109+
110+
def test_establish_connection_using_3_levels_config_with_integer_keys
111+
previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
112+
default_shard_was = ActiveRecord::Base.default_shard
113+
ActiveRecord::Base.default_shard = 0
114+
115+
config = {
116+
"default_env" => {
117+
"primary" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3" },
118+
"primary_shard_one" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_one.sqlite3" },
119+
"primary_shard_one_replica" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_one_replica.sqlite3", "replica" => true }
120+
}
121+
}
122+
123+
@prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
124+
125+
IntegerKeysBase.connects_to(shards: {
126+
0 => { writing: :primary, reading: :primary },
127+
1 => { writing: :primary_shard_one, reading: :primary_shard_one_replica }
128+
})
129+
130+
connection_description_name = "ActiveRecord::ConnectionAdapters::ConnectionHandlersShardingDbTest::IntegerKeysBase"
131+
base_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_description_name)
132+
default_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_description_name, shard: 0)
133+
134+
assert_equal [0, 1], ActiveRecord::Base.connection_handler.send(:get_pool_manager, connection_description_name).shard_names
135+
assert_equal base_pool, default_pool
136+
assert_equal "test/db/primary.sqlite3", default_pool.db_config.database
137+
assert_equal "primary", default_pool.db_config.name
138+
139+
assert_not_nil pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_description_name, shard: 1)
140+
assert_equal "test/db/primary_shard_one.sqlite3", pool.db_config.database
141+
assert_equal "primary_shard_one", pool.db_config.name
142+
143+
assert_not_nil pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_description_name, role: :reading, shard: 1)
144+
assert_equal "test/db/primary_shard_one_replica.sqlite3", pool.db_config.database
145+
assert_equal "primary_shard_one_replica", pool.db_config.name
146+
147+
IntegerKeysBase.connected_to(shard: 1) do
148+
assert_includes IntegerKeysBase.lease_connection.query_value("SELECT file FROM pragma_database_list;"), "primary_shard_one.sqlite3"
149+
end
150+
151+
IntegerKeysBase.connected_to(role: :reading, shard: 1) do
152+
assert_includes IntegerKeysBase.lease_connection.query_value("SELECT file FROM pragma_database_list;"), "primary_shard_one_replica.sqlite3"
153+
end
154+
155+
assert_raises(ActiveRecord::ConnectionNotDefined) do
156+
IntegerKeysBase.connected_to(shard: 2) do
157+
IntegerKeysBase.lease_connection.query_value("SELECT file FROM pragma_database_list;")
158+
end
159+
end
160+
ensure
161+
ActiveRecord::Base.default_shard = default_shard_was
162+
ActiveRecord::Base.configurations = @prev_configs
163+
ActiveRecord::Base.establish_connection(:arunit)
164+
ENV["RAILS_ENV"] = previous_env
165+
end
166+
106167
def test_switching_connections_via_handler
107168
previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
108169

0 commit comments

Comments
 (0)