Move terminate_all_connections into setup Rake task

1 files changed, 21 insertions, 1 deletions

diff --git a/lib/tasks/gitlab/setup.rake b/lib/tasks/gitlab/setup.rake
index dab96c5c54d..e876b23d43f 100644
--- a/lib/tasks/gitlab/setup.rake
+++ b/lib/tasks/gitlab/setup.rake
@@ -28,7 +28,7 @@ namespace :gitlab do
     # In production, we might want to prevent ourselves from shooting
     # ourselves in the foot, so let's only do this in a test or
     # development environment.
-    Gitlab::Database.terminate_all_connections unless Rails.env.production?
+    terminate_all_connections unless Rails.env.production?
 
     Rake::Task["db:reset"].invoke
     Rake::Task["add_limits_mysql"].invoke
@@ -38,4 +38,24 @@ namespace :gitlab do
     puts "Quitting...".color(:red)
     exit 1
   end
+
+  # If there are any clients connected to the DB, PostgreSQL won't let
+  # you drop the database. It's possible that Sidekiq, Unicorn, or
+  # some other client will be hanging onto a connection, preventing
+  # the DROP DATABASE from working. To workaround this problem, this
+  # method terminates all the connections so that a subsequent DROP
+  # will work.
+  def self.terminate_all_connections
+    return false unless Gitlab::Database.postgresql?
+
+    cmd = <<~SQL
+    SELECT pg_terminate_backend(pg_stat_activity.pid)
+        FROM pg_stat_activity
+        WHERE datname = current_database()
+        AND pid <> pg_backend_pid();
+    SQL
+
+    ActiveRecord::Base.connection.execute(cmd)&.result_status == PG::PGRES_TUPLES_OK
+  rescue ActiveRecord::NoDatabaseError
+  end
 end