From 5488e8a4a3a2bb1493a58429d53e32c94f18114e Mon Sep 17 00:00:00 2001
From: Stan Hu <stanhu@gmail.com>
Date: Mon, 28 Jan 2019 14:46:05 -0800
Subject: Kill all PostgreSQL connections for rake dev:setup

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 cancels all the
connections so that the db:reset command will work.

Note that there's still a slight possibility a client connects after its
connection is terminated. If this is an issue, we could solve it by
revoking CONNECT access, but for now it seems this works.

Closes https://gitlab.com/gitlab-org/gitlab-development-kit/issues/450
---
 lib/gitlab/database.rb           | 22 ++++++++++++++++++++++
 lib/tasks/gitlab/setup.rake      |  5 +++++
 spec/lib/gitlab/database_spec.rb | 12 ++++++++++++
 3 files changed, 39 insertions(+)

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index b6ca777e029..817b5b97a3b 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -235,6 +235,28 @@ module Gitlab
       connection.schema_cache.data_source_exists?(table_name)
     end
 
+    # WARNING: Only run this if you know what you're doing.
+    #
+    # 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
+
+      connection.execute(cmd)&.result_status == PG::PGRES_TUPLES_OK
+    rescue ActiveRecord::NoDatabaseError
+    end
+
     private_class_method :connection
 
     def self.database_version
diff --git a/lib/tasks/gitlab/setup.rake b/lib/tasks/gitlab/setup.rake
index f71e69987cb..dab96c5c54d 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.
+    Gitlab::Database.terminate_all_connections unless Rails.env.production?
+
     Rake::Task["db:reset"].invoke
     Rake::Task["add_limits_mysql"].invoke
     Rake::Task["setup_postgresql"].invoke
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index 60106ee3c0b..cbc2ce719be 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -410,6 +410,18 @@ describe Gitlab::Database do
     end
   end
 
+  describe '.terminate_all_connections' do
+    it 'returns true for PostgreSQL', :postgresql do
+      expect(described_class.terminate_all_connections).to be_truthy
+    end
+
+    it 'returns false for MySQL' do
+      expect(described_class).to receive(:postgresql?).and_return(false)
+
+      expect(described_class.terminate_all_connections).to be_falsey
+    end
+  end
+
   describe '#true_value' do
     it 'returns correct value for PostgreSQL' do
       expect(described_class).to receive(:postgresql?).and_return(true)
-- 
cgit v1.2.1