diff --git a/activerecord/lib/active_record/connection_adapters.rb b/activerecord/lib/active_record/connection_adapters.rb index 0d4dc941fc..d7626150c9 100644 --- a/activerecord/lib/active_record/connection_adapters.rb +++ b/activerecord/lib/active_record/connection_adapters.rb @@ -10,6 +10,7 @@ module ActiveRecord autoload :Column autoload :Role + autoload :RoleManager autoload :Resolver autoload_at "active_record/connection_adapters/abstract/schema_definitions" do diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb index 41041e9548..3b5c23d2fb 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb @@ -1004,7 +1004,7 @@ module ActiveRecord def initialize # These caches are keyed by role.connection_specification_name (Role#connection_specification_name). - @owner_to_role = Concurrent::Map.new(initial_capacity: 2) + @owner_to_role_manager = Concurrent::Map.new(initial_capacity: 2) # Backup finalizer: if the forked child skipped Kernel#fork the early discard has not occurred ObjectSpace.define_finalizer self, FINALIZER @@ -1031,20 +1031,20 @@ module ActiveRecord end def connection_pool_names # :nodoc: - owner_to_role.keys + owner_to_role_manager.keys end def connection_pool_list - owner_to_role.values.compact.map(&:pool) + owner_to_role_manager.values.compact.flat_map { |rc| rc.roles.map(&:pool) } end alias :connection_pools :connection_pool_list - def establish_connection(config) + def establish_connection(config, role_name = :default) resolver = Resolver.new(Base.configurations) role = resolver.resolve_role(config) db_config = role.db_config - remove_connection(role.connection_specification_name) + remove_connection(role.connection_specification_name, role_name) message_bus = ActiveSupport::Notifications.instrumenter payload = { @@ -1055,7 +1055,9 @@ module ActiveRecord payload[:config] = db_config.configuration_hash end - owner_to_role[role.connection_specification_name] = role + owner_to_role_manager[role.connection_specification_name] ||= RoleManager.new + role_manager = owner_to_role_manager[role.connection_specification_name] + role_manager.set_role(role_name, role) message_bus.instrument("!connection.active_record", payload) do role.pool @@ -1114,8 +1116,8 @@ module ActiveRecord # Returns true if a connection that's accessible to this class has # already been opened. - def connected?(spec_name) - pool = retrieve_connection_pool(spec_name) + def connected?(spec_name, role_name = :default) + pool = retrieve_connection_pool(spec_name, role_name) pool && pool.connected? end @@ -1123,22 +1125,27 @@ module ActiveRecord # connection and the defined connection (if they exist). The result # can be used as an argument for #establish_connection, for easily # re-establishing the connection. - def remove_connection(spec_name) - if role = owner_to_role.delete(spec_name) - role.disconnect! - role.db_config.configuration_hash + def remove_connection(spec_name, role_name = :default) + if role_manager = owner_to_role_manager[spec_name] + role = role_manager.remove_role(role_name) + + if role + role.disconnect! + role.db_config.configuration_hash + end end end - # Retrieving the connection pool happens a lot, so we cache it in @owner_to_role. + # Retrieving the connection pool happens a lot, so we cache it in @owner_to_role_manager. # This makes retrieving the connection pool O(1) once the process is warm. # When a connection is established or removed, we invalidate the cache. - def retrieve_connection_pool(spec_name) - owner_to_role[spec_name]&.pool + def retrieve_connection_pool(spec_name, role_name = :default) + role = owner_to_role_manager[spec_name]&.get_role(role_name) + role&.pool end private - attr_reader :owner_to_role + attr_reader :owner_to_role_manager end end end diff --git a/activerecord/lib/active_record/connection_adapters/role_manager.rb b/activerecord/lib/active_record/connection_adapters/role_manager.rb new file mode 100644 index 0000000000..3837cdbb0e --- /dev/null +++ b/activerecord/lib/active_record/connection_adapters/role_manager.rb @@ -0,0 +1,27 @@ +# frozen_string_literal: true + +module ActiveRecord + module ConnectionAdapters + class RoleManager # :nodoc: + def initialize + @name_to_role = {} + end + + def roles + @name_to_role.values + end + + def remove_role(name) + @name_to_role.delete(name) + end + + def get_role(name) + @name_to_role[name] + end + + def set_role(name, role) + @name_to_role[name] = role + end + end + end +end diff --git a/activerecord/lib/active_record/test_fixtures.rb b/activerecord/lib/active_record/test_fixtures.rb index 60c9a84121..30d3804ce3 100644 --- a/activerecord/lib/active_record/test_fixtures.rb +++ b/activerecord/lib/active_record/test_fixtures.rb @@ -192,7 +192,11 @@ module ActiveRecord ActiveRecord::Base.connection_handlers.values.each do |handler| if handler != writing_handler handler.connection_pool_names.each do |name| - handler.send(:owner_to_role)[name] = writing_handler.send(:owner_to_role)[name] + writing_role_manager = writing_handler.send(:owner_to_role_manager)[name] + writing_role = writing_role_manager.get_role(:default) + + role_manager = handler.send(:owner_to_role_manager)[name] + role_manager.set_role(:default, writing_role) end end end diff --git a/activerecord/test/cases/connection_adapters/connection_handlers_multi_role_test.rb b/activerecord/test/cases/connection_adapters/connection_handlers_multi_role_test.rb new file mode 100644 index 0000000000..292472fc36 --- /dev/null +++ b/activerecord/test/cases/connection_adapters/connection_handlers_multi_role_test.rb @@ -0,0 +1,106 @@ +# frozen_string_literal: true + +require "cases/helper" +require "models/person" + +module ActiveRecord + module ConnectionAdapters + class ConnectionHandlersMultiRoleTest < ActiveRecord::TestCase + self.use_transactional_tests = false + + fixtures :people + + def setup + @handlers = { writing: ConnectionHandler.new } + @rw_handler = @handlers[:writing] + @spec_name = "primary" + @writing_handler = ActiveRecord::Base.connection_handlers[:writing] + end + + def teardown + ActiveRecord::Base.connection_handlers = { writing: ActiveRecord::Base.default_connection_handler } + end + + unless in_memory_db? + def test_establish_connection_with_roles + previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env" + + config = { + "default_env" => { + "primary" => { "adapter" => "sqlite3", "database" => "db/primary.sqlite3" } + } + } + + @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config + + @writing_handler.establish_connection(:primary) + @writing_handler.establish_connection(:primary, :role_two) + + default_pool = @writing_handler.retrieve_connection_pool("primary", :default) + other_pool = @writing_handler.retrieve_connection_pool("primary", :role_two) + + assert_not_nil default_pool + assert_not_equal default_pool, other_pool + + # :default if passed with no key + assert_equal default_pool, @writing_handler.retrieve_connection_pool("primary") + ensure + ActiveRecord::Base.configurations = @prev_configs + ActiveRecord::Base.establish_connection(:arunit) + ENV["RAILS_ENV"] = previous_env + end + + def test_remove_connection + previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env" + + config = { + "default_env" => { + "primary" => { "adapter" => "sqlite3", "database" => "db/primary.sqlite3" } + } + } + + @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config + + @writing_handler.establish_connection(:primary) + @writing_handler.establish_connection(:primary, :role_two) + + # remove default + @writing_handler.remove_connection("primary") + + assert_nil @writing_handler.retrieve_connection_pool("primary") + assert_not_nil @writing_handler.retrieve_connection_pool("primary", :role_two) + ensure + ActiveRecord::Base.configurations = @prev_configs + ActiveRecord::Base.establish_connection(:arunit) + ENV["RAILS_ENV"] = previous_env + end + + def test_connected? + previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env" + + config = { + "default_env" => { + "primary" => { "adapter" => "sqlite3", "database" => "db/primary.sqlite3" } + } + } + + @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config + + @writing_handler.establish_connection(:primary) + @writing_handler.establish_connection(:primary, :role_two) + + # connect to default + @writing_handler.connection_pool_list.first.checkout + + assert @writing_handler.connected?("primary") + assert @writing_handler.connected?("primary", :default) + assert_not @writing_handler.connected?("primary", :role_two) + ensure + ActiveRecord::Base.configurations = @prev_configs + ActiveRecord::Base.establish_connection(:arunit) + ENV["RAILS_ENV"] = previous_env + end + end + end + end +end diff --git a/activerecord/test/cases/fixtures_test.rb b/activerecord/test/cases/fixtures_test.rb index ce48560c0d..9c0e9e5ea6 100644 --- a/activerecord/test/cases/fixtures_test.rb +++ b/activerecord/test/cases/fixtures_test.rb @@ -1414,7 +1414,7 @@ class MultipleDatabaseFixturesTest < ActiveRecord::TestCase private def with_temporary_connection_pool - role = ActiveRecord::Base.connection_handler.send(:owner_to_role).fetch("primary") + role = ActiveRecord::Base.connection_handler.send(:owner_to_role_manager).fetch("primary").get_role(:default) new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(role) role.stub(:pool, new_pool) do diff --git a/activerecord/test/cases/query_cache_test.rb b/activerecord/test/cases/query_cache_test.rb index c418816877..e35bc7010b 100644 --- a/activerecord/test/cases/query_cache_test.rb +++ b/activerecord/test/cases/query_cache_test.rb @@ -555,7 +555,7 @@ class QueryCacheTest < ActiveRecord::TestCase private def with_temporary_connection_pool - role = ActiveRecord::Base.connection_handler.send(:owner_to_role).fetch("primary") + role = ActiveRecord::Base.connection_handler.send(:owner_to_role_manager).fetch("primary").get_role(:default) new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(role) role.stub(:pool, new_pool) do