1 files changed, 24 insertions, 14 deletions
diff --git a/spec/support/db_cleaner.rb b/spec/support/db_cleaner.rb
index e3a05f17593..24cdbe04fc2 100644
--- a/spec/support/db_cleaner.rb
+++ b/spec/support/db_cleaner.rb
@@ -78,22 +78,32 @@ module DbCleaner
     puts "Databases re-creation done in #{Gitlab::Metrics::System.monotonic_time - start}"
   end
 
+  def recreate_databases_and_seed_if_needed
+    # Postgres maximum number of columns in a table is 1600 (https://github.com/postgres/postgres/blob/de41869b64d57160f58852eab20a27f248188135/src/include/access/htup_details.h#L23-L47).
+    # We drop and recreate the database if any table has more than 1200 columns, just to be safe.
+    return false unless any_connection_class_with_more_than_allowed_columns?
+
+    recreate_all_databases!
+
+    # Seed required data as recreating DBs will delete it
+    TestEnv.seed_db
+
+    true
+  end
+
   def force_disconnect_all_connections!
-    all_connection_classes.each do |connection_class|
-      # We use `connection_pool` to avoid going through
-      # Load Balancer since it does retry ops
-      pool = connection_class.connection_pool
-
-      # Force disconnect https://www.cybertec-postgresql.com/en/terminating-database-connections-in-postgresql/
-      pool.connection.execute(<<-SQL)
-        SELECT pg_terminate_backend(pid)
-          FROM pg_stat_activity
-          WHERE datname = #{pool.connection.quote(pool.db_config.database)}
-            AND pid != pg_backend_pid();
-      SQL
-
-      connection_class.connection_pool.disconnect!
+    cmd = <<~SQL
+      SELECT pg_terminate_backend(pg_stat_activity.pid)
+      FROM pg_stat_activity
+      WHERE datname = current_database()
+        AND pid <> pg_backend_pid();
+    SQL
+
+    Gitlab::Database::EachDatabase.each_database_connection(include_shared: false) do |connection|
+      connection.execute(cmd)
     end
+
+    ActiveRecord::Base.clear_all_connections! # rubocop:disable Database/MultipleDatabases
   end
 end