1 files changed, 25 insertions, 0 deletions
diff --git a/lib/tasks/gitlab/setup.rake b/lib/tasks/gitlab/setup.rake
index f71e69987cb..e876b23d43f 100644
--- a/lib/tasks/gitlab/setup.rake
+++ b/lib/tasks/gitlab/setup.rake
@@ -25,6 +25,11 @@ namespace :gitlab do
       puts ""
     end
 
+    # 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.
+    terminate_all_connections unless Rails.env.production?
+
     Rake::Task["db:reset"].invoke
     Rake::Task["add_limits_mysql"].invoke
     Rake::Task["setup_postgresql"].invoke
@@ -33,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