From 061b527d79020fd8bc9cb163a88a88dbedfe3d4c Mon Sep 17 00:00:00 2001 From: eileencodes Date: Tue, 23 Nov 2021 09:50:59 -0500 Subject: [PATCH] Expose role/shard on pool/connection It can be useful to know what the role and shard of the connection or pool is. Previously there was no way to find the role or shard other than asking `connected_to?` - sometimes it's better to just have the actual role and shard names available. --- .../abstract/connection_handler.rb | 6 ++-- .../abstract/connection_pool.rb | 4 ++- .../connection_adapters/abstract_adapter.rb | 12 ++++++++ .../connection_adapters/pool_config.rb | 6 ++-- .../adapter_leasing_test.rb | 2 +- .../test/cases/connection_pool_test.rb | 30 ++++++++++++++++--- activerecord/test/cases/reaper_test.rb | 4 +-- 7 files changed, 51 insertions(+), 13 deletions(-) diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_handler.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_handler.rb index bab5e2a08d..68a305fb12 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/connection_handler.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_handler.rb @@ -126,7 +126,7 @@ module ActiveRecord def establish_connection(config, owner_name: Base, role: ActiveRecord::Base.current_role, shard: Base.current_shard) owner_name = StringConnectionOwner.new(config.to_s) if config.is_a?(Symbol) - pool_config = resolve_pool_config(config, owner_name) + pool_config = resolve_pool_config(config, owner_name, role, shard) db_config = pool_config.db_config # Protects the connection named `ActiveRecord::Base` from being removed @@ -264,7 +264,7 @@ module ActiveRecord # pool_config.db_config.configuration_hash # # => { host: "localhost", database: "foo", adapter: "sqlite3" } # - def resolve_pool_config(config, owner_name) + def resolve_pool_config(config, owner_name, role, shard) db_config = Base.configurations.resolve(config) raise(AdapterNotSpecified, "database configuration does not specify adapter") unless db_config.adapter @@ -294,7 +294,7 @@ module ActiveRecord raise AdapterNotFound, "database configuration specifies nonexistent #{db_config.adapter} adapter" end - ConnectionAdapters::PoolConfig.new(owner_name, db_config) + ConnectionAdapters::PoolConfig.new(owner_name, db_config, role, shard) end end end 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 221f3786f2..9ae3b331de 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb @@ -105,7 +105,7 @@ module ActiveRecord include ConnectionAdapters::AbstractPool attr_accessor :automatic_reconnect, :checkout_timeout - attr_reader :db_config, :size, :reaper, :pool_config, :connection_klass, :async_executor + attr_reader :db_config, :size, :reaper, :pool_config, :connection_klass, :async_executor, :role, :shard delegate :schema_cache, :schema_cache=, to: :pool_config @@ -121,6 +121,8 @@ module ActiveRecord @pool_config = pool_config @db_config = pool_config.db_config @connection_klass = pool_config.connection_klass + @role = pool_config.role + @shard = pool_config.shard @checkout_timeout = db_config.checkout_timeout @idle_timeout = db_config.idle_timeout diff --git a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb index 48c7a0fa9d..094881e6e5 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb +++ b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb @@ -224,6 +224,18 @@ module ActiveRecord @pool.connection_klass end + # The role (ie :writing) for the current connection. In a + # non-multi role application, `:writing` is returned. + def role + @pool.role + end + + # The shard (ie :default) for the current connection. In + # a non-sharded application, `:default` is returned. + def shard + @pool.shard + end + def schema_cache @pool.get_schema_cache(self) end diff --git a/activerecord/lib/active_record/connection_adapters/pool_config.rb b/activerecord/lib/active_record/connection_adapters/pool_config.rb index ff048f15be..c3ab9b4c38 100644 --- a/activerecord/lib/active_record/connection_adapters/pool_config.rb +++ b/activerecord/lib/active_record/connection_adapters/pool_config.rb @@ -5,7 +5,7 @@ module ActiveRecord class PoolConfig # :nodoc: include Mutex_m - attr_reader :db_config, :connection_klass + attr_reader :db_config, :connection_klass, :role, :shard attr_accessor :schema_cache INSTANCES = ObjectSpace::WeakMap.new @@ -17,10 +17,12 @@ module ActiveRecord end end - def initialize(connection_klass, db_config) + def initialize(connection_klass, db_config, role, shard) super() @connection_klass = connection_klass @db_config = db_config + @role = role + @shard = shard @pool = nil INSTANCES[self] = self end diff --git a/activerecord/test/cases/connection_adapters/adapter_leasing_test.rb b/activerecord/test/cases/connection_adapters/adapter_leasing_test.rb index 06337f601b..8652763275 100644 --- a/activerecord/test/cases/connection_adapters/adapter_leasing_test.rb +++ b/activerecord/test/cases/connection_adapters/adapter_leasing_test.rb @@ -40,7 +40,7 @@ module ActiveRecord def test_close db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new("test", "primary", {}) - pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config) + pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config, :writing, :default) pool = Pool.new(pool_config) pool.insert_connection_for_test! @adapter @adapter.pool = pool diff --git a/activerecord/test/cases/connection_pool_test.rb b/activerecord/test/cases/connection_pool_test.rb index 53b0924cd4..37f1eb756f 100644 --- a/activerecord/test/cases/connection_pool_test.rb +++ b/activerecord/test/cases/connection_pool_test.rb @@ -13,7 +13,7 @@ module ActiveRecord # Keep a duplicate pool so we do not bother others @db_config = ActiveRecord::Base.connection_pool.db_config - @pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, @db_config) + @pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, @db_config, :writing, :default) @pool = ConnectionPool.new(@pool_config) if in_memory_db? @@ -204,7 +204,7 @@ module ActiveRecord config = @db_config.configuration_hash.merge(idle_timeout: "0.02") db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(@db_config.env_name, @db_config.name, config) - pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config) + pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config, :writing, :default) @pool = ConnectionPool.new(pool_config) idle_conn = @pool.checkout @pool.checkin(idle_conn) @@ -231,7 +231,7 @@ module ActiveRecord config = @db_config.configuration_hash.merge(idle_timeout: -5) db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new(@db_config.env_name, @db_config.name, config) - pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config) + pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config, :writing, :default) @pool = ConnectionPool.new(pool_config) idle_conn = @pool.checkout @pool.checkin(idle_conn) @@ -739,11 +739,33 @@ module ActiveRecord assert_not_nil found_conn end + def test_role_and_shard_is_returned + assert_equal :writing, @pool_config.role + assert_equal :writing, @pool.role + assert_equal :writing, @pool.connection.role + + assert_equal :default, @pool_config.shard + assert_equal :default, @pool.shard + assert_equal :default, @pool.connection.shard + + db_config = ActiveRecord::Base.connection_pool.db_config + pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config, :reading, :shard_one) + pool = ConnectionPool.new(pool_config) + + assert_equal :reading, pool_config.role + assert_equal :reading, pool.role + assert_equal :reading, pool.connection.role + + assert_equal :shard_one, pool_config.shard + assert_equal :shard_one, pool.shard + assert_equal :shard_one, pool.connection.shard + end + private def with_single_connection_pool config = @db_config.configuration_hash.merge(pool: 1) db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new("arunit", "primary", config) - pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config) + pool_config = ActiveRecord::ConnectionAdapters::PoolConfig.new(ActiveRecord::Base, db_config, :writing, :default) yield(pool = ConnectionPool.new(pool_config)) ensure diff --git a/activerecord/test/cases/reaper_test.rb b/activerecord/test/cases/reaper_test.rb index 8c79fd577b..4f71557d99 100644 --- a/activerecord/test/cases/reaper_test.rb +++ b/activerecord/test/cases/reaper_test.rb @@ -59,7 +59,7 @@ module ActiveRecord def test_pool_has_reaper config = ActiveRecord::Base.configurations.configs_for(env_name: "arunit", name: "primary") - pool_config = PoolConfig.new(ActiveRecord::Base, config) + pool_config = PoolConfig.new(ActiveRecord::Base, config, :writing, :default) pool = ConnectionPool.new(pool_config) assert pool.reaper @@ -171,7 +171,7 @@ module ActiveRecord def duplicated_pool_config(merge_config_options = {}) old_config = ActiveRecord::Base.connection_pool.db_config.configuration_hash.merge(merge_config_options) db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new("arunit", "primary", old_config.dup) - PoolConfig.new(ActiveRecord::Base, db_config) + PoolConfig.new(ActiveRecord::Base, db_config, :writing, :default) end def new_conn_in_thread(pool)