From ba2f38e498179baca89fad865630e5c16fc375d7 Mon Sep 17 00:00:00 2001 From: eileencodes Date: Fri, 14 Aug 2020 11:05:05 -0400 Subject: [PATCH] Fix incorrect removal of current_shard in establish_connection If we enter a `connected_to` block and call `establish_connection` like the test added here we need to ensure that `shard: current_shard` is passed to the handler, otherwise the connection will be established on `default` not on `shard_one`. Co-authored-by: John Crepezzi --- activerecord/lib/active_record/connection_handling.rb | 2 +- .../connection_handlers_sharding_db_test.rb | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/activerecord/lib/active_record/connection_handling.rb b/activerecord/lib/active_record/connection_handling.rb index cb9ca865af..8ddf8dc183 100644 --- a/activerecord/lib/active_record/connection_handling.rb +++ b/activerecord/lib/active_record/connection_handling.rb @@ -49,7 +49,7 @@ module ActiveRecord def establish_connection(config_or_env = nil) config_or_env ||= DEFAULT_ENV.call.to_sym db_config, owner_name = resolve_config_for_connection(config_or_env) - connection_handler.establish_connection(db_config, owner_name: owner_name) + connection_handler.establish_connection(db_config, owner_name: owner_name, shard: current_shard) end # Connects a model to the databases specified. The +database+ keyword diff --git a/activerecord/test/cases/connection_adapters/connection_handlers_sharding_db_test.rb b/activerecord/test/cases/connection_adapters/connection_handlers_sharding_db_test.rb index 5536d51acc..6d0e17320a 100644 --- a/activerecord/test/cases/connection_adapters/connection_handlers_sharding_db_test.rb +++ b/activerecord/test/cases/connection_adapters/connection_handlers_sharding_db_test.rb @@ -25,6 +25,16 @@ module ActiveRecord end unless in_memory_db? + def test_establishing_a_connection_in_connected_to_block_uses_current_role_and_shard + ActiveRecord::Base.connected_to(shard: :shard_one) do + db_config = ActiveRecord::Base.configurations.configs_for(env_name: "arunit", name: "primary") + ActiveRecord::Base.establish_connection(db_config) + assert_nothing_raised { Person.first } + + assert_equal [:default, :shard_one], ActiveRecord::Base.connection_handlers[:writing].send(:owner_to_pool_manager).fetch("ActiveRecord::Base").instance_variable_get(:@name_to_pool_config).keys + end + end + def test_establish_connection_using_3_levels_config previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"