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.
This commit is contained in:
eileencodes 2021-11-23 09:50:59 -05:00
parent b5946d5df7
commit 061b527d79
No known key found for this signature in database
GPG Key ID: BA5C575120BBE8DF
7 changed files with 51 additions and 13 deletions

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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

View File

@ -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)