Merge pull request rails#55680 from HeyNonster/nony/integer-database-config-keys

Support integer shard keys

Author: byroot

activerecord/CHANGELOG.md

@@ -1,3 +1,18 @@
+*   Add support for integer shard keys.
+    ```ruby
+    # Now accepts symbols as shard keys.
+    ActiveRecord::Base.connects_to(shards: {
+      1: { writing: :primary_shard_one, reading: :primary_shard_one },
+      2: { writing: :primary_shard_two, reading: :primary_shard_two},
+    })
+
+    ActiveRecord::Base.connected_to(shard: 1) do
+      # ..
+    end
+    ```
+
+    *Nony Dutton*
+
 *   Add `ActiveRecord::Base.only_columns`
 
     Similar in use case to `ignored_columns` but listing columns to consider rather than the ones

activerecord/lib/active_record/connection_handling.rb

@@ -100,7 +100,8 @@ def connects_to(database: {}, shards: {})
           db_config = resolve_config_for_connection(database_key)
 
           self.connection_class = true
-          connections << connection_handler.establish_connection(db_config, owner_name: self, role: role, shard: shard.to_sym)
+          shard = shard.to_sym unless shard.is_a? Integer
+          connections << connection_handler.establish_connection(db_config, owner_name: self, role: role, shard: shard)
         end
       end
 

activerecord/test/cases/connection_adapters/connection_handlers_sharding_db_test.rb

@@ -103,6 +103,67 @@ def test_establish_connection_using_3_levels_config_with_shards_and_replica
           ENV["RAILS_ENV"] = previous_env
         end
 
+        class IntegerKeysBase < ActiveRecord::Base
+          self.abstract_class = true
+        end
+
+        def test_establish_connection_using_3_levels_config_with_integer_keys
+          previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
+          default_shard_was = ActiveRecord::Base.default_shard
+          ActiveRecord::Base.default_shard = 0
+
+          config = {
+            "default_env" => {
+              "primary" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3" },
+              "primary_shard_one" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_one.sqlite3" },
+              "primary_shard_one_replica" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_one_replica.sqlite3", "replica" => true }
+            }
+          }
+
+          @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
+
+          IntegerKeysBase.connects_to(shards: {
+            0 => { writing: :primary, reading: :primary },
+            1 => { writing: :primary_shard_one, reading: :primary_shard_one_replica }
+          })
+
+          connection_description_name = "ActiveRecord::ConnectionAdapters::ConnectionHandlersShardingDbTest::IntegerKeysBase"
+          base_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_description_name)
+          default_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_description_name, shard: 0)
+
+          assert_equal [0, 1], ActiveRecord::Base.connection_handler.send(:get_pool_manager, connection_description_name).shard_names
+          assert_equal base_pool, default_pool
+          assert_equal "test/db/primary.sqlite3", default_pool.db_config.database
+          assert_equal "primary", default_pool.db_config.name
+
+          assert_not_nil pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_description_name, shard: 1)
+          assert_equal "test/db/primary_shard_one.sqlite3", pool.db_config.database
+          assert_equal "primary_shard_one", pool.db_config.name
+
+          assert_not_nil pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool(connection_description_name, role: :reading, shard: 1)
+          assert_equal "test/db/primary_shard_one_replica.sqlite3", pool.db_config.database
+          assert_equal "primary_shard_one_replica", pool.db_config.name
+
+          IntegerKeysBase.connected_to(shard: 1) do
+            assert_includes IntegerKeysBase.lease_connection.query_value("SELECT file FROM pragma_database_list;"), "primary_shard_one.sqlite3"
+          end
+
+          IntegerKeysBase.connected_to(role: :reading, shard: 1) do
+            assert_includes IntegerKeysBase.lease_connection.query_value("SELECT file FROM pragma_database_list;"), "primary_shard_one_replica.sqlite3"
+          end
+
+          assert_raises(ActiveRecord::ConnectionNotDefined) do
+            IntegerKeysBase.connected_to(shard: 2) do
+              IntegerKeysBase.lease_connection.query_value("SELECT file FROM pragma_database_list;")
+            end
+          end
+        ensure
+          ActiveRecord::Base.default_shard = default_shard_was
+          ActiveRecord::Base.configurations = @prev_configs
+          ActiveRecord::Base.establish_connection(:arunit)
+          ENV["RAILS_ENV"] = previous_env
+        end
+
         def test_switching_connections_via_handler
           previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"