fix(referential-integrity): make sure fks are added back

This would silently fail in the test
`test_bulk_insert_with_a_multi_statement_query_raises_an_exception_when_any_insert_fails`.
Removing some foreign keys and thus failing afterwards
in the test suite. Moreover, the `#disable_referential_integrity`
method is unfortunately public, so we need it to be robust.

Author: BuonOmo

lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb

@@ -35,8 +35,30 @@ def check_all_foreign_keys_valid!
 
         def disable_referential_integrity
           foreign_keys = all_foreign_keys
-          dropped_fks = foreign_keys
 
+          remove_foreign_keys(foreign_keys)
+
+          # Prefixes and suffixes are added in add_foreign_key
+          # in AR7+ so we need to temporarily disable them here,
+          # otherwise prefixes/suffixes will be erroneously added.
+          old_prefix = ActiveRecord::Base.table_name_prefix
+          old_suffix = ActiveRecord::Base.table_name_suffix
+
+          begin
+            yield
+          ensure
+            ActiveRecord::Base.table_name_prefix = ""
+            ActiveRecord::Base.table_name_suffix = ""
+            add_foreign_keys(foreign_keys) # Never raises.
+
+            ActiveRecord::Base.table_name_prefix = old_prefix if defined?(old_prefix)
+            ActiveRecord::Base.table_name_suffix = old_suffix if defined?(old_suffix)
+          end
+        end
+
+        private
+
+        def remove_foreign_keys(foreign_keys)
           statements = foreign_keys.map do |foreign_key|
             # We do not use the `#remove_foreign_key` method here because it
             # checks for foreign keys existance in the schema cache. This method
@@ -47,54 +69,28 @@ def disable_referential_integrity
             schema_creation.accept(at)
           end
           execute_batch(statements, "Disable referential integrity -> remove foreign keys")
+        end
 
-          STDOUT.puts "BEFORE FK"
-          yield
-          STDOUT.puts "AFTER FK"
-
-          # Prefixes and suffixes are added in add_foreign_key
-          # in AR7+ so we need to temporarily disable them here,
-          # otherwise prefixes/suffixes will be erroneously added.
-          old_prefix = ActiveRecord::Base.table_name_prefix
-          old_suffix = ActiveRecord::Base.table_name_suffix
+        # NOTE: This method should never raise, otherwise we risk polluting table name
+        #   prefixes and suffixes. The good thing is: if this happens, tests will crash
+        #   hard, no way we miss it.
+        def add_foreign_keys(foreign_keys)
+          # We avoid using `foreign_key_exists?` here because it checks the schema cache
+          # for every key. This method is performance critical for the test suite, hence
+          # we use the `#all_foreign_keys` method that only make one query to the database.
+          already_inserted_foreign_keys = all_foreign_keys
+          statements = foreign_keys.map do |foreign_key|
+            next if already_inserted_foreign_keys.any? { |fk| fk.from_table == foreign_key.from_table && fk.options[:name] == foreign_key.options[:name] }
 
-          ActiveRecord::Base.table_name_prefix = ""
-          ActiveRecord::Base.table_name_suffix = ""
+            options = foreign_key_options(foreign_key.from_table, foreign_key.to_table, foreign_key.options)
+            at = create_alter_table foreign_key.from_table
+            at.add_foreign_key foreign_key.to_table, options
 
-          begin
-            # Avoid having PG:DuplicateObject error if a test is ran in transaction.
-            # TODO: verify that there is no cache issue related to running this (e.g: fk
-            #   still in cache but not in db)
-            #
-            # We avoid using `foreign_key_exists?` here because it checks the schema cache
-            # for every key. This method is performance critical for the test suite, hence
-            # we use the `#all_foreign_keys` method that only make one query to the database.
-            already_inserted_foreign_keys = all_foreign_keys
-            new_fks = []
-            statements = foreign_keys.map do |foreign_key|
-              next if already_inserted_foreign_keys.any? { |fk| fk.from_table == foreign_key.from_table && fk.options[:name] == foreign_key.options[:name] }
-
-              options = foreign_key_options(foreign_key.from_table, foreign_key.to_table, foreign_key.options)
-              at = create_alter_table foreign_key.from_table
-              at.add_foreign_key foreign_key.to_table, options
-
-              new_fks << foreign_key
-              schema_creation.accept(at)
-            end
-            execute_batch(statements.compact, "Disable referential integrity -> add foreign keys")
-
-            new_fks = new_fks.map { _1.from_table + " -> " + _1.to_table + (_1.options[:name] ? " (#{_1.options[:name]})" : "") }
-            dropped_fks = dropped_fks.map { _1.from_table + " -> " + _1.to_table + (_1.options[:name] ? " (#{_1.options[:name]})" : "") }
-            STDOUT.puts "FK DIFF NEW:\n", (new_fks - dropped_fks).map{ "\t#{_1}" } if (new_fks - dropped_fks).any?
-            STDOUT.puts "FK DIFF OLD:\n", (dropped_fks - new_fks).map{ "\t#{_1}" } if (dropped_fks - new_fks).any?
-          ensure
-            ActiveRecord::Base.table_name_prefix = old_prefix
-            ActiveRecord::Base.table_name_suffix = old_suffix
+            schema_creation.accept(at)
           end
+          execute_batch(statements.compact, "Disable referential integrity -> add foreign keys")
         end
 
-        private
-
         # NOTE: Copy/paste of the `#foreign_keys(table)` method adapted
         #   to return every single foreign key in the database.
         def all_foreign_keys