Merge pull request rails#51010 from Shopify/handle-transactional-fixtures-leak

Gracefully handle transactional fixtures leaks

Author: byroot

activerecord/lib/active_record/test_fixtures.rb

@@ -130,25 +130,53 @@ def setup_fixtures(config = ActiveRecord::Base)
       end
 
       @fixture_cache = {}
-      @fixture_connections = []
+      @fixture_connections = {}
       @fixture_cache_key = [self.class.fixture_table_names.dup, self.class.fixture_paths.dup, self.class.fixture_class_names.dup]
       @@already_loaded_fixtures ||= {}
       @connection_subscriber = nil
       @saved_pool_configs = Hash.new { |hash, key| hash[key] = {} }
 
-      # Load fixtures once and begin transaction.
       if run_in_transaction?
+        # Load fixtures once and begin transaction.
         @loaded_fixtures = @@already_loaded_fixtures[@fixture_cache_key]
         unless @loaded_fixtures
           @@already_loaded_fixtures.clear
           @loaded_fixtures = @@already_loaded_fixtures[@fixture_cache_key] = load_fixtures(config)
         end
 
+        setup_transactional_fixtures
+      else
+        # Load fixtures for every test.
+        ActiveRecord::FixtureSet.reset_cache
+        invalidate_already_loaded_fixtures
+        @loaded_fixtures = load_fixtures(config)
+      end
+
+      # Instantiate fixtures for every test if requested.
+      instantiate_fixtures if use_instantiated_fixtures
+    end
+
+    def teardown_fixtures
+      # Rollback changes if a transaction is active.
+      if run_in_transaction?
+        teardown_transactional_fixtures
+      else
+        ActiveRecord::FixtureSet.reset_cache
+      end
+
+      ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
+    end
+
+    private
+      def setup_transactional_fixtures
+        setup_shared_connection_pool
+
         # Begin transactions for connections already established
-        @fixture_connections = enlist_fixture_connections
-        @fixture_connections.each do |connection|
-          connection.begin_transaction joinable: false, _lazy: false
-          connection.pool.lock_thread = true if lock_threads
+        @fixture_connections = ActiveRecord::Base.connection_handler.connection_pool_list(:writing).to_h do |pool|
+          pool.lock_thread = true if lock_threads
+          connection = pool.connection
+          transaction = connection.begin_transaction(joinable: false, _lazy: false)
+          [connection, transaction]
         end
 
         # When connections are established in the future, begin a transaction too
@@ -167,50 +195,38 @@ def setup_fixtures(config = ActiveRecord::Base)
             if connection
               setup_shared_connection_pool
 
-              if !@fixture_connections.include?(connection)
-                connection.begin_transaction joinable: false, _lazy: false
+              if !@fixture_connections.key?(connection)
                 connection.pool.lock_thread = true if lock_threads
-                @fixture_connections << connection
+                @fixture_connections[connection] = connection.begin_transaction(joinable: false, _lazy: false)
               end
             end
           end
         end
-
-      # Load fixtures for every test.
-      else
-        ActiveRecord::FixtureSet.reset_cache
-        @@already_loaded_fixtures.clear
-        @loaded_fixtures = load_fixtures(config)
       end
 
-      # Instantiate fixtures for every test if requested.
-      instantiate_fixtures if use_instantiated_fixtures
-    end
-
-    def teardown_fixtures
-      # Rollback changes if a transaction is active.
-      if run_in_transaction?
+      def teardown_transactional_fixtures
         ActiveSupport::Notifications.unsubscribe(@connection_subscriber) if @connection_subscriber
-        @fixture_connections.each do |connection|
-          connection.rollback_transaction if connection.transaction_open?
+        @fixture_connections.each do |connection, transaction|
+          begin
+            connection.rollback_transaction(transaction)
+          rescue ActiveRecord::StatementInvalid
+            # Something commited or rolled back the transaction.
+            # We can no longer trust the database state is clean.
+            invalidate_already_loaded_fixtures
+            # We also don't know for sure the connection wasn't
+            # mutated in dangerous ways.
+            connection.disconnect!
+          end
           connection.pool.lock_thread = false
         end
         @fixture_connections.clear
         teardown_shared_connection_pool
-      else
-        ActiveRecord::FixtureSet.reset_cache
       end
 
-      ActiveRecord::Base.connection_handler.clear_active_connections!(:all)
-    end
-
-    def enlist_fixture_connections
-      setup_shared_connection_pool
-
-      ActiveRecord::Base.connection_handler.connection_pool_list(:writing).map(&:connection)
-    end
+      def invalidate_already_loaded_fixtures
+        @@already_loaded_fixtures.clear
+      end
 
-    private
       # Shares the writing connection pool with connections on
       # other handlers.
       #

activerecord/test/cases/active_record_test.rb

@@ -4,9 +4,11 @@
 require "rack"
 
 class ActiveRecordTest < ActiveRecord::TestCase
+  self.use_transactional_tests = false
+
   unless in_memory_db?
     test ".disconnect_all! closes all connections" do
-      ActiveRecord::Base.connection.active?
+      ActiveRecord::Base.connection.connect!
       assert_predicate ActiveRecord::Base, :connected?
 
       ActiveRecord.disconnect_all!

activerecord/test/cases/adapters/abstract_mysql_adapter/connection_test.rb

@@ -37,7 +37,7 @@ def test_no_automatic_reconnection_after_timeout
     assert_not_predicate @connection, :active?
   ensure
     # Repair all fixture connections so other tests won't break.
-    @fixture_connections.each(&:verify!)
+    @fixture_connections.each_key(&:verify!)
   end
 
   def test_successful_reconnection_after_timeout_with_manual_reconnect

activerecord/test/cases/adapters/postgresql/connection_test.rb

@@ -145,7 +145,7 @@ def test_reconnection_after_actual_disconnection_with_verify
       assert_predicate @connection, :active?
     ensure
       # Repair all fixture connections so other tests won't break.
-      @fixture_connections.each(&:verify!)
+      @fixture_connections.each_key(&:verify!)
     end
 
     def test_set_session_variable_true

activerecord/test/cases/connection_adapters/connection_handler_test.rb

@@ -6,6 +6,8 @@
 module ActiveRecord
   module ConnectionAdapters
     class ConnectionHandlerTest < ActiveRecord::TestCase
+      self.use_transactional_tests = false
+
       fixtures :people
 
       def setup
@@ -93,6 +95,8 @@ def test_fixtures_dont_raise_if_theres_no_writing_pool_config
           connection_handler = ActiveRecord::Base.connection_handler
           ActiveRecord::Base.connection_handler = ActiveRecord::ConnectionAdapters::ConnectionHandler.new
 
+          setup_transactional_fixtures
+
           assert_nothing_raised do
             ActiveRecord::Base.connects_to(database: { reading: :arunit, writing: :arunit })
           end
@@ -101,6 +105,8 @@ def test_fixtures_dont_raise_if_theres_no_writing_pool_config
           ro_conn = ActiveRecord::Base.connection_handler.retrieve_connection("ActiveRecord::Base", role: :reading)
 
           assert_equal rw_conn, ro_conn
+
+          teardown_transactional_fixtures
         ensure
           ActiveRecord::Base.connection_handler = connection_handler
         end

activerecord/test/cases/connection_adapters/schema_cache_test.rb

@@ -5,6 +5,8 @@
 module ActiveRecord
   module ConnectionAdapters
     class SchemaCacheTest < ActiveRecord::TestCase
+      self.use_transactional_tests = false
+
       def setup
         @connection = ARUnit2Model.connection
         @cache = new_bound_reflection

activerecord/test/cases/primary_class_test.rb

@@ -3,6 +3,8 @@
 require "cases/helper"
 
 class PrimaryClassTest < ActiveRecord::TestCase
+  self.use_transactional_tests = false
+
   def teardown
     clean_up_connection_handler
   end