Merge pull request #40384 from eugeneius/teardown_shared_connection_pool

Restore connection pools after transactional tests

activerecord/lib/active_record/test_fixtures.rb

@@ -114,6 +114,9 @@ def setup_fixtures(config = ActiveRecord::Base)
 
       # Load fixtures once and begin transaction.
       if run_in_transaction?
+        @legacy_saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
+        @saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
+
         if @@already_loaded_fixtures[self.class]
           @loaded_fixtures = @@already_loaded_fixtures[self.class]
         else
@@ -169,6 +172,7 @@ def teardown_fixtures
           connection.pool.lock_thread = false
         end
         @fixture_connections.clear
+        teardown_shared_connection_pool
       else
         ActiveRecord::FixtureSet.reset_cache
       end
@@ -200,8 +204,13 @@ def setup_shared_connection_pool
                 return unless writing_pool_manager
 
                 pool_manager = handler.send(:owner_to_pool_manager)[name]
+                @legacy_saved_pool_configs[handler][name] ||= {}
                 pool_manager.shard_names.each do |shard_name|
                   writing_pool_config = writing_pool_manager.get_pool_config(nil, shard_name)
+                  pool_config = pool_manager.get_pool_config(nil, shard_name)
+                  next if pool_config == writing_pool_config
+
+                  @legacy_saved_pool_configs[handler][name][shard_name] = pool_config
                   pool_manager.set_pool_config(nil, shard_name, writing_pool_config)
                 end
               end
@@ -214,15 +223,48 @@ def setup_shared_connection_pool
             pool_manager = handler.send(:owner_to_pool_manager)[name]
             pool_manager.shard_names.each do |shard_name|
               writing_pool_config = pool_manager.get_pool_config(ActiveRecord::Base.writing_role, shard_name)
+              @saved_pool_configs[name][shard_name] ||= {}
               pool_manager.role_names.each do |role|
-                next unless pool_manager.get_pool_config(role, shard_name)
+                next unless pool_config = pool_manager.get_pool_config(role, shard_name)
+                next if pool_config == writing_pool_config
+
+                @saved_pool_configs[name][shard_name][role] = pool_config
                 pool_manager.set_pool_config(role, shard_name, writing_pool_config)
               end
             end
           end
         end
       end
 
+      def teardown_shared_connection_pool
+        if ActiveRecord::Base.legacy_connection_handling
+          @legacy_saved_pool_configs.each_pair do |handler, names|
+            names.each_pair do |name, shards|
+              shards.each_pair do |shard_name, pool_config|
+                pool_manager = handler.send(:owner_to_pool_manager)[name]
+                pool_manager.set_pool_config(nil, shard_name, pool_config)
+              end
+            end
+          end
+        else
+          handler = ActiveRecord::Base.connection_handler
+
+          @saved_pool_configs.each_pair do |name, shards|
+            pool_manager = handler.send(:owner_to_pool_manager)[name]
+            shards.each_pair do |shard_name, roles|
+              roles.each_pair do |role, pool_config|
+                next unless pool_manager.get_pool_config(role, shard_name)
+
+                pool_manager.set_pool_config(role, shard_name, pool_config)
+              end
+            end
+          end
+        end
+
+        @legacy_saved_pool_configs.clear
+        @saved_pool_configs.clear
+      end
+
       def load_fixtures(config)
         ActiveRecord::FixtureSet.create_fixtures(fixture_path, fixture_table_names, fixture_class_names, config).index_by(&:name)
       end

activerecord/test/cases/fixtures_test.rb

@@ -1452,6 +1452,13 @@ def test_writing_and_reading_connections_are_the_same
       ro_conn = handler.retrieve_connection_pool("ActiveRecord::Base", role: :reading).connection
 
       assert_equal rw_conn, ro_conn
+
+      teardown_shared_connection_pool
+
+      rw_conn = handler.retrieve_connection_pool("ActiveRecord::Base", role: :writing).connection
+      ro_conn = handler.retrieve_connection_pool("ActiveRecord::Base", role: :reading).connection
+
+      assert_not_equal rw_conn, ro_conn
     end
 
     def test_writing_and_reading_connections_are_the_same_for_non_default_shards
@@ -1465,6 +1472,13 @@ def test_writing_and_reading_connections_are_the_same_for_non_default_shards
       ro_conn = handler.retrieve_connection_pool("ActiveRecord::Base", role: :reading, shard: :two).connection
 
       assert_equal rw_conn, ro_conn
+
+      teardown_shared_connection_pool
+
+      rw_conn = handler.retrieve_connection_pool("ActiveRecord::Base", role: :writing, shard: :two).connection
+      ro_conn = handler.retrieve_connection_pool("ActiveRecord::Base", role: :reading, shard: :two).connection
+
+      assert_not_equal rw_conn, ro_conn
     end
 
     def test_only_existing_connections_are_replaced
@@ -1482,6 +1496,17 @@ def test_only_existing_connections_are_replaced
       end
     end
 
+    def test_only_existing_connections_are_restored
+      clean_up_connection_handler
+      teardown_shared_connection_pool
+
+      assert_raises(ActiveRecord::ConnectionNotEstablished) do
+        ActiveRecord::Base.connected_to(role: :reading) do
+          ActiveRecord::Base.retrieve_connection
+        end
+      end
+    end
+
     private
       def config
         { "default" => default_config, "readonly" => readonly_config }
@@ -1543,6 +1568,13 @@ def test_writing_and_reading_connections_are_the_same_with_legacy_handling
       ro_conn = reading.retrieve_connection_pool("ActiveRecord::Base").connection
 
       assert_equal rw_conn, ro_conn
+
+      teardown_shared_connection_pool
+
+      rw_conn = writing.retrieve_connection_pool("ActiveRecord::Base").connection
+      ro_conn = reading.retrieve_connection_pool("ActiveRecord::Base").connection
+
+      assert_not_equal rw_conn, ro_conn
     end
 
     def test_writing_and_reading_connections_are_the_same_for_non_default_shards_with_legacy_handling
@@ -1558,6 +1590,13 @@ def test_writing_and_reading_connections_are_the_same_for_non_default_shards_wit
       ro_conn = reading.retrieve_connection_pool("ActiveRecord::Base", shard: :two).connection
 
       assert_equal rw_conn, ro_conn
+
+      teardown_shared_connection_pool
+
+      rw_conn = writing.retrieve_connection_pool("ActiveRecord::Base", shard: :two).connection
+      ro_conn = reading.retrieve_connection_pool("ActiveRecord::Base", shard: :two).connection
+
+      assert_not_equal rw_conn, ro_conn
     end
 
     def test_only_existing_connections_are_replaced
@@ -1575,6 +1614,17 @@ def test_only_existing_connections_are_replaced
       end
     end
 
+    def test_only_existing_connections_are_restored
+      clean_up_legacy_connection_handlers
+      teardown_shared_connection_pool
+
+      assert_raises(ActiveRecord::ConnectionNotEstablished) do
+        ActiveRecord::Base.connected_to(role: :reading) do
+          ActiveRecord::Base.retrieve_connection
+        end
+      end
+    end
+
     private
       def config
         { "default" => default_config, "readonly" => readonly_config }