Move terminate_all_connections into setup Rake task

3 files changed, 21 insertions, 35 deletions

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 817b5b97a3b..b6ca777e029 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -235,28 +235,6 @@ 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 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
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index cbc2ce719be..60106ee3c0b 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -410,18 +410,6 @@ 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)