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.

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

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

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

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

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

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

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)