Merge pull request rails#54366 from flavorjones/flavorjones-dont-cache-class-in-pool-config-v3

`PoolConfig` no longer keeps a reference to the connection class (attempt 3)

Author: byroot

activerecord/CHANGELOG.md

@@ -1,3 +1,10 @@
+*   PoolConfig no longer keeps a reference to the connection class.
+
+    Keeping a reference to the class caused subtle issues when combined with reloading in
+    development. Fixes #54343.
+
+    *Mike Dalessio*
+
 *   Fix SQL notifications sometimes not sent when using async queries.
 
     ```ruby

activerecord/lib/active_record/connection_adapters/abstract/connection_handler.rb

@@ -54,19 +54,22 @@ module ConnectionAdapters
     # about the model. The model needs to pass a connection specification name to the handler,
     # in order to look up the correct connection pool.
     class ConnectionHandler
-      class StringConnectionName # :nodoc:
-        attr_reader :name
-
-        def initialize(name)
+      class ConnectionDescriptor # :nodoc:
+        def initialize(name, primary = false)
           @name = name
+          @primary = primary
+        end
+
+        def name
+          primary_class? ? "ActiveRecord::Base" : @name
         end
 
         def primary_class?
-          false
+          @primary
         end
 
         def current_preventing_writes
-          false
+          ActiveRecord::Base.preventing_writes?(@name)
         end
       end
 
@@ -115,7 +118,7 @@ def establish_connection(config, owner_name: Base, role: Base.current_role, shar
         pool_config = resolve_pool_config(config, owner_name, role, shard)
         db_config = pool_config.db_config
 
-        pool_manager = set_pool_manager(pool_config.connection_name)
+        pool_manager = set_pool_manager(pool_config.connection_descriptor)
 
         # If there is an existing pool with the same values as the pool_config
         # don't remove the connection. Connections should only be removed if we are
@@ -127,8 +130,8 @@ def establish_connection(config, owner_name: Base, role: Base.current_role, shar
           # Update the pool_config's connection class if it differs. This is used
           # for ensuring that ActiveRecord::Base and the primary_abstract_class use
           # the same pool. Without this granular swapping will not work correctly.
-          if owner_name.primary_class? && (existing_pool_config.connection_class != owner_name)
-            existing_pool_config.connection_class = owner_name
+          if owner_name.primary_class? && (existing_pool_config.connection_descriptor != owner_name)
+            existing_pool_config.connection_descriptor = owner_name
           end
 
           existing_pool_config.pool
@@ -137,7 +140,7 @@ def establish_connection(config, owner_name: Base, role: Base.current_role, shar
           pool_manager.set_pool_config(role, shard, pool_config)
 
           payload = {
-            connection_name: pool_config.connection_name,
+            connection_name: pool_config.connection_descriptor.name,
             role: role,
             shard: shard,
             config: db_config.configuration_hash
@@ -242,8 +245,8 @@ def get_pool_manager(connection_name)
         end
 
         # Get the existing pool manager or initialize and assign a new one.
-        def set_pool_manager(connection_name)
-          connection_name_to_pool_manager[connection_name] ||= PoolManager.new
+        def set_pool_manager(connection_descriptor)
+          connection_name_to_pool_manager[connection_descriptor.name] ||= PoolManager.new
         end
 
         def pool_managers
@@ -278,9 +281,9 @@ def resolve_pool_config(config, connection_name, role, shard)
 
         def determine_owner_name(owner_name, config)
           if owner_name.is_a?(String) || owner_name.is_a?(Symbol)
-            StringConnectionName.new(owner_name.to_s)
+            ConnectionDescriptor.new(owner_name.to_s)
           elsif config.is_a?(Symbol)
-            StringConnectionName.new(config.to_s)
+            ConnectionDescriptor.new(config.to_s)
           else
             owner_name
           end

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -36,7 +36,7 @@ def schema_reflection
       end
 
       def schema_cache; end
-      def connection_class; end
+      def connection_descriptor; end
       def checkin(_); end
       def remove(_); end
       def async_executor; end
@@ -364,8 +364,8 @@ def unpin_connection! # :nodoc:
         clean
       end
 
-      def connection_class # :nodoc:
-        pool_config.connection_class
+      def connection_descriptor # :nodoc:
+        pool_config.connection_descriptor
       end
 
       # Returns true if there is an open connection being used for the current thread.

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -226,9 +226,9 @@ def default_timezone
       # the value of +current_preventing_writes+.
       def preventing_writes?
         return true if replica?
-        return false if connection_class.nil?
+        return false if connection_descriptor.nil?
 
-        connection_class.current_preventing_writes
+        connection_descriptor.current_preventing_writes
       end
 
       def prepared_statements?
@@ -279,8 +279,8 @@ def lease
         @owner = ActiveSupport::IsolatedExecutionState.context
       end
 
-      def connection_class # :nodoc:
-        @pool.connection_class
+      def connection_descriptor # :nodoc:
+        @pool.connection_descriptor
       end
 
       # The role (e.g. +:writing+) for the current connection. In a

activerecord/lib/active_record/connection_adapters/pool_config.rb

@@ -5,9 +5,8 @@ module ConnectionAdapters
     class PoolConfig # :nodoc:
       include MonitorMixin
 
-      attr_reader :db_config, :role, :shard
+      attr_reader :db_config, :role, :shard, :connection_descriptor
       attr_writer :schema_reflection, :server_version
-      attr_accessor :connection_class
 
       def schema_reflection
         @schema_reflection ||= SchemaReflection.new(db_config.lazy_schema_cache_path)
@@ -29,7 +28,7 @@ def disconnect_all!
       def initialize(connection_class, db_config, role, shard)
         super()
         @server_version = nil
-        @connection_class = connection_class
+        self.connection_descriptor = connection_class
         @db_config = db_config
         @role = role
         @shard = shard
@@ -41,11 +40,12 @@ def server_version(connection)
         @server_version || synchronize { @server_version ||= connection.get_database_version }
       end
 
-      def connection_name
-        if connection_class.primary_class?
-          "ActiveRecord::Base"
+      def connection_descriptor=(connection_descriptor)
+        case connection_descriptor
+        when ConnectionHandler::ConnectionDescriptor
+          @connection_descriptor = connection_descriptor
         else
-          connection_class.name
+          @connection_descriptor = ConnectionHandler::ConnectionDescriptor.new(connection_descriptor.name, connection_descriptor.primary_class?)
         end
       end
 

activerecord/lib/active_record/core.rb

@@ -201,6 +201,17 @@ def self.current_preventing_writes
         false
       end
 
+      # Intended to behave like `.current_preventing_writes` given the class name as input.
+      # See PoolConfig and ConnectionHandler::ConnectionDescriptor.
+      def self.preventing_writes?(class_name) # :nodoc:
+        connected_to_stack.reverse_each do |hash|
+          return hash[:prevent_writes] if !hash[:prevent_writes].nil? && hash[:klasses].include?(Base)
+          return hash[:prevent_writes] if !hash[:prevent_writes].nil? && hash[:klasses].any? { |klass| klass.name == class_name }
+        end
+
+        false
+      end
+
       def self.connected_to_stack # :nodoc:
         if connected_to_stack = ActiveSupport::IsolatedExecutionState[:active_record_connected_to_stack]
           connected_to_stack

activerecord/test/cases/connection_adapters/connection_swapping_nested_test.rb

@@ -460,6 +460,34 @@ def test_application_record_prevent_writes_can_be_changed
           Object.send(:remove_const, :ApplicationRecord)
           ActiveRecord::Base.establish_connection :arunit
         end
+
+        def test_prevent_writes_handles_class_reloading
+          # Regression test for https://github.com/rails/rails/issues/54343
+          Object.const_set(:ReloadedRecord, Class.new(ActiveRecord::Base) { self.abstract_class = true })
+          ReloadedRecord.connects_to(database: { writing: :arunit, reading: :arunit })
+
+          ActiveRecord::Base.connected_to(role: :reading, prevent_writes: true) do
+            ReloadedRecord.connected_to(role: :writing, prevent_writes: false) do
+              assert_not_predicate ReloadedRecord.lease_connection, :preventing_writes?
+            end
+          end
+
+          # emulate a reload in development mode
+          Object.send(:remove_const, :ReloadedRecord)
+          Object.const_set(:ReloadedRecord, Class.new(ActiveRecord::Base) { self.abstract_class = true })
+          ReloadedRecord.connects_to(database: { writing: :arunit, reading: :arunit })
+
+          ActiveRecord::Base.connected_to(role: :reading, prevent_writes: true) do
+            ReloadedRecord.connected_to(role: :writing, prevent_writes: false) do
+              assert_not_predicate ReloadedRecord.lease_connection, :preventing_writes?
+            end
+          end
+        ensure
+          ReloadedRecord.remove_connection
+          ActiveRecord.application_record_class = nil
+          Object.send(:remove_const, :ReloadedRecord)
+          ActiveRecord::Base.establish_connection :arunit
+        end
       end
     end
   end

activerecord/test/cases/multi_db_migrator_test.rb

@@ -71,8 +71,8 @@ def puts(*)
   def test_schema_migration_is_different_for_different_connections
     assert_not_equal @schema_migration_a, @schema_migration_b
     assert_not_equal @schema_migration_a.instance_variable_get(:@pool), @schema_migration_b.instance_variable_get(:@pool)
-    assert_equal "ActiveRecord::Base", @pool_a.pool_config.connection_name
-    assert_equal "ARUnit2Model", @pool_b.pool_config.connection_name
+    assert_equal "ActiveRecord::Base", @pool_a.pool_config.connection_descriptor.name
+    assert_equal "ARUnit2Model", @pool_b.pool_config.connection_descriptor.name
   end
 
   def test_finds_migrations