Add latest changes from gitlab-org/gitlab@14-10-stable-eev14.10.0-rc42

6 files changed, 777 insertions, 164 deletions

diff --git a/spec/tasks/dev_rake_spec.rb b/spec/tasks/dev_rake_spec.rb
index 7bc27d2732c..73b1604aa10 100644
--- a/spec/tasks/dev_rake_spec.rb
+++ b/spec/tasks/dev_rake_spec.rb
@@ -7,9 +7,20 @@ RSpec.describe 'dev rake tasks' do
     Rake.application.rake_require 'tasks/gitlab/setup'
     Rake.application.rake_require 'tasks/gitlab/shell'
     Rake.application.rake_require 'tasks/dev'
+    Rake.application.rake_require 'active_record/railties/databases'
+    Rake.application.rake_require 'tasks/gitlab/db'
   end
 
   describe 'setup' do
+    around do |example|
+      old_force_value = ENV['force']
+
+      # setup rake task sets the force env var, so reset it
+      example.run
+
+      ENV['force'] = old_force_value # rubocop:disable RSpec/EnvAssignment
+    end
+
     subject(:setup_task) { run_rake_task('dev:setup') }
 
     let(:connections) { Gitlab::Database.database_base_models.values.map(&:connection) }
@@ -17,7 +28,9 @@ RSpec.describe 'dev rake tasks' do
     it 'sets up the development environment', :aggregate_failures do
       expect(Rake::Task['gitlab:setup']).to receive(:invoke)
 
+      expect(connections).to all(receive(:execute).with('SET statement_timeout TO 0'))
       expect(connections).to all(receive(:execute).with('ANALYZE'))
+      expect(connections).to all(receive(:execute).with('RESET statement_timeout'))
 
       expect(Rake::Task['gitlab:shell:setup']).to receive(:invoke)
 
@@ -35,4 +48,103 @@ RSpec.describe 'dev rake tasks' do
       load_task
     end
   end
+
+  describe 'terminate_all_connections' do
+    let(:connections) do
+      Gitlab::Database.database_base_models.values.filter_map do |model|
+        model.connection if Gitlab::Database.db_config_share_with(model.connection_db_config).nil?
+      end
+    end
+
+    def expect_connections_to_be_terminated
+      expect(Gitlab::Database::EachDatabase).to receive(:each_database_connection)
+        .with(include_shared: false)
+        .and_call_original
+
+      expect(connections).to all(receive(:execute).with(/SELECT pg_terminate_backend/))
+    end
+
+    def expect_connections_not_to_be_terminated
+      connections.each do |connection|
+        expect(connection).not_to receive(:execute)
+      end
+    end
+
+    subject(:terminate_task) { run_rake_task('dev:terminate_all_connections') }
+
+    it 'terminates all connections' do
+      expect_connections_to_be_terminated
+
+      terminate_task
+    end
+
+    context 'when in the production environment' do
+      it 'does not terminate connections' do
+        expect(Rails.env).to receive(:production?).and_return(true)
+        expect_connections_not_to_be_terminated
+
+        terminate_task
+      end
+    end
+
+    context 'when a database is not found' do
+      before do
+        skip_if_multiple_databases_not_setup
+      end
+
+      it 'continues to next connection' do
+        expect(connections.first).to receive(:execute).and_raise(ActiveRecord::NoDatabaseError)
+        expect(connections.second).to receive(:execute).with(/SELECT pg_terminate_backend/)
+
+        terminate_task
+      end
+    end
+  end
+
+  context 'multiple databases' do
+    before do
+      skip_if_multiple_databases_not_setup
+    end
+
+    context 'with a valid database' do
+      describe 'copy_db:ci' do
+        before do
+          allow(Rake::Task['dev:terminate_all_connections']).to receive(:invoke)
+
+          configurations = instance_double(ActiveRecord::DatabaseConfigurations)
+          allow(ActiveRecord::Base).to receive(:configurations).and_return(configurations)
+          allow(configurations).to receive(:configs_for).with(env_name: Rails.env, name: 'ci').and_return(ci_configuration)
+        end
+
+        subject(:load_task) { run_rake_task('dev:setup_ci_db') }
+
+        let(:ci_configuration) { instance_double(ActiveRecord::DatabaseConfigurations::HashConfig, name: 'ci', database: '__test_db_ci') }
+
+        it 'creates the database from main' do
+          expect(ApplicationRecord.connection).to receive(:create_database).with(
+            ci_configuration.database,
+            template: ApplicationRecord.connection_db_config.database
+          )
+
+          expect(Rake::Task['dev:terminate_all_connections']).to receive(:invoke)
+
+          run_rake_task('dev:copy_db:ci')
+        end
+
+        context 'when the database already exists' do
+          it 'prints out a warning' do
+            expect(ApplicationRecord.connection).to receive(:create_database).and_raise(ActiveRecord::DatabaseAlreadyExists)
+
+            expect { run_rake_task('dev:copy_db:ci') }.to output(/Database '#{ci_configuration.database}' already exists/).to_stderr
+          end
+        end
+      end
+    end
+
+    context 'with an invalid database' do
+      it 'raises an error' do
+        expect { run_rake_task('dev:copy_db:foo') }.to raise_error(RuntimeError, /Don't know how to build task/)
+      end
+    end
+  end
 end
diff --git a/spec/tasks/gitlab/backup_rake_spec.rb b/spec/tasks/gitlab/backup_rake_spec.rb
index df9f2a0d3bb..6080948403d 100644
--- a/spec/tasks/gitlab/backup_rake_spec.rb
+++ b/spec/tasks/gitlab/backup_rake_spec.rb
@@ -199,18 +199,25 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         end
 
         it 'logs the progress to log file' do
-          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping database ... ")
-          expect(Gitlab::BackupLogger).to receive(:info).with(message: "[SKIPPED]")
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping database ... [SKIPPED]")
           expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping repositories ... ")
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping repositories ... done")
           expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping uploads ... ")
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping uploads ... done")
           expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping builds ... ")
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping builds ... done")
           expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping artifacts ... ")
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping artifacts ... done")
           expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping pages ... ")
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping pages ... done")
           expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping lfs objects ... ")
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping lfs objects ... done")
           expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping terraform states ... ")
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping terraform states ... done")
           expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping container registry images ... ")
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping container registry images ... done")
           expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping packages ... ")
-          expect(Gitlab::BackupLogger).to receive(:info).with(message: "done").exactly(9).times
+          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping packages ... done")
 
           backup_tasks.each do |task|
             run_rake_task("gitlab:backup:#{task}:create")
@@ -228,19 +235,19 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
       db_backup_error = Backup::DatabaseBackupError.new(config, db_file_name)
 
       where(:backup_class, :rake_task, :error) do
-        Backup::Database     | 'gitlab:backup:db:create'        | db_backup_error
-        Backup::Builds       | 'gitlab:backup:builds:create'    | file_backup_error
-        Backup::Uploads      | 'gitlab:backup:uploads:create'   | file_backup_error
-        Backup::Artifacts    | 'gitlab:backup:artifacts:create' | file_backup_error
-        Backup::Pages        | 'gitlab:backup:pages:create'     | file_backup_error
-        Backup::Lfs          | 'gitlab:backup:lfs:create'       | file_backup_error
-        Backup::Registry     | 'gitlab:backup:registry:create'  | file_backup_error
+        Backup::Database | 'gitlab:backup:db:create'        | db_backup_error
+        Backup::Files    | 'gitlab:backup:builds:create'    | file_backup_error
+        Backup::Files    | 'gitlab:backup:uploads:create'   | file_backup_error
+        Backup::Files    | 'gitlab:backup:artifacts:create' | file_backup_error
+        Backup::Files    | 'gitlab:backup:pages:create'     | file_backup_error
+        Backup::Files    | 'gitlab:backup:lfs:create'       | file_backup_error
+        Backup::Files    | 'gitlab:backup:registry:create'  | file_backup_error
       end
 
       with_them do
         before do
-          expect_next_instance_of(backup_class) do |instance|
-            expect(instance).to receive(:dump).and_raise(error)
+          allow_next_instance_of(backup_class) do |instance|
+            allow(instance).to receive(:dump).and_raise(error)
           end
         end
 
@@ -408,25 +415,12 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         create(:project, :repository)
       end
 
-      it 'has defaults' do
-        expect(::Backup::Repositories).to receive(:new)
-          .with(anything, strategy: anything, max_concurrency: 1, max_storage_concurrency: 1)
-          .and_call_original
-
-        expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
-      end
-
       it 'passes through concurrency environment variables' do
-        # The way concurrency is handled will change with the `gitaly_backup`
-        # feature flag. For now we need to check that both ways continue to
-        # work. This will be cleaned up in the rollout issue.
-        # See https://gitlab.com/gitlab-org/gitlab/-/issues/333034
-
         stub_env('GITLAB_BACKUP_MAX_CONCURRENCY', 5)
         stub_env('GITLAB_BACKUP_MAX_STORAGE_CONCURRENCY', 2)
 
         expect(::Backup::Repositories).to receive(:new)
-          .with(anything, strategy: anything, max_concurrency: 5, max_storage_concurrency: 2)
+          .with(anything, strategy: anything)
           .and_call_original
         expect(::Backup::GitalyBackup).to receive(:new).with(anything, max_parallelism: 5, storage_parallelism: 2, incremental: false).and_call_original
 
diff --git a/spec/tasks/gitlab/db/validate_config_rake_spec.rb b/spec/tasks/gitlab/db/validate_config_rake_spec.rb
new file mode 100644
index 00000000000..0b2c844a91f
--- /dev/null
+++ b/spec/tasks/gitlab/db/validate_config_rake_spec.rb
@@ -0,0 +1,205 @@
+# frozen_string_literal: true
+
+require 'rake_helper'
+
+RSpec.describe 'gitlab:db:validate_config', :silence_stdout do
+  before :all do
+    Rake.application.rake_require 'active_record/railties/databases'
+    Rake.application.rake_require 'tasks/seed_fu'
+    Rake.application.rake_require 'tasks/gitlab/db/validate_config'
+
+    # empty task as env is already loaded
+    Rake::Task.define_task :environment
+  end
+
+  context "when validating config" do
+    let(:main_database_config) do
+      Rails.application.config.load_database_yaml
+        .dig('test', 'main')
+        .slice('adapter', 'encoding', 'database', 'username', 'password', 'host')
+        .symbolize_keys
+    end
+
+    let(:additional_database_config) do
+      # Use built-in postgres database
+      main_database_config.merge(database: 'postgres')
+    end
+
+    around do |example|
+      with_reestablished_active_record_base(reconnect: true) do
+        with_db_configs(test: test_config) do
+          example.run
+        end
+      end
+    end
+
+    shared_examples 'validates successfully' do
+      it 'by default' do
+        expect { run_rake_task('gitlab:db:validate_config') }.not_to output(/Database config validation failure/).to_stderr
+        expect { run_rake_task('gitlab:db:validate_config') }.not_to raise_error
+      end
+
+      it 'for production' do
+        allow(Gitlab).to receive(:dev_or_test_env?).and_return(false)
+
+        expect { run_rake_task('gitlab:db:validate_config') }.not_to output(/Database config validation failure/).to_stderr
+        expect { run_rake_task('gitlab:db:validate_config') }.not_to raise_error
+      end
+
+      it 'always re-establishes ActiveRecord::Base connection to main config' do
+        run_rake_task('gitlab:db:validate_config')
+
+        expect(ActiveRecord::Base.connection_db_config.configuration_hash).to include(main_database_config) # rubocop: disable Database/MultipleDatabases
+      end
+
+      it 'if GITLAB_VALIDATE_DATABASE_CONFIG is set' do
+        stub_env('GITLAB_VALIDATE_DATABASE_CONFIG', '1')
+        allow(Gitlab).to receive(:dev_or_test_env?).and_return(false)
+
+        expect { run_rake_task('gitlab:db:validate_config') }.not_to output(/Database config validation failure/).to_stderr
+        expect { run_rake_task('gitlab:db:validate_config') }.not_to raise_error
+      end
+
+      context 'when finding the initializer fails' do
+        where(:raised_error) { [ActiveRecord::NoDatabaseError, ActiveRecord::ConnectionNotEstablished, PG::ConnectionBad] }
+        with_them do
+          it "does not raise an error for #{params[:raised_error]}" do
+            allow(ActiveRecord::Base.connection).to receive(:select_one).and_raise(raised_error) # rubocop: disable Database/MultipleDatabases
+
+            expect { run_rake_task('gitlab:db:validate_config') }.not_to output(/Database config validation failure/).to_stderr
+            expect { run_rake_task('gitlab:db:validate_config') }.not_to raise_error
+          end
+        end
+      end
+    end
+
+    shared_examples 'raises an error' do |match|
+      it 'by default' do
+        expect { run_rake_task('gitlab:db:validate_config') }.to raise_error(match)
+      end
+
+      it 'for production' do
+        allow(Gitlab).to receive(:dev_or_test_env?).and_return(false)
+
+        expect { run_rake_task('gitlab:db:validate_config') }.to raise_error(match)
+      end
+
+      it 'always re-establishes ActiveRecord::Base connection to main config' do
+        expect { run_rake_task('gitlab:db:validate_config') }.to raise_error(match)
+
+        expect(ActiveRecord::Base.connection_db_config.configuration_hash).to include(main_database_config) # rubocop: disable Database/MultipleDatabases
+      end
+
+      it 'if GITLAB_VALIDATE_DATABASE_CONFIG=1' do
+        stub_env('GITLAB_VALIDATE_DATABASE_CONFIG', '1')
+
+        expect { run_rake_task('gitlab:db:validate_config') }.to raise_error(match)
+      end
+
+      it 'to stderr if GITLAB_VALIDATE_DATABASE_CONFIG=0' do
+        stub_env('GITLAB_VALIDATE_DATABASE_CONFIG', '0')
+
+        expect { run_rake_task('gitlab:db:validate_config') }.to output(match).to_stderr
+      end
+    end
+
+    context 'when only main: is specified' do
+      let(:test_config) do
+        {
+          main: main_database_config
+        }
+      end
+
+      it_behaves_like 'validates successfully'
+    end
+
+    context 'when main: uses database_tasks=false' do
+      let(:test_config) do
+        {
+          main: main_database_config.merge(database_tasks: false)
+        }
+      end
+
+      it_behaves_like 'raises an error', /The 'main' is required to use 'database_tasks: true'/
+    end
+
+    context 'when many configurations share the same database' do
+      context 'when no database_tasks is specified, assumes true' do
+        let(:test_config) do
+          {
+            main: main_database_config,
+            ci: main_database_config
+          }
+        end
+
+        it_behaves_like 'raises an error', /Many configurations \(main, ci\) share the same database/
+      end
+
+      context 'when database_tasks is specified' do
+        let(:test_config) do
+          {
+            main: main_database_config.merge(database_tasks: true),
+            ci: main_database_config.merge(database_tasks: true)
+          }
+        end
+
+        it_behaves_like 'raises an error', /Many configurations \(main, ci\) share the same database/
+      end
+
+      context "when there's no main: but something different, as currently we only can share with main:" do
+        let(:test_config) do
+          {
+            archive: main_database_config,
+            ci: main_database_config.merge(database_tasks: false)
+          }
+        end
+
+        it_behaves_like 'raises an error', /The 'ci' is expecting to share configuration with 'main', but no such is to be found/
+      end
+    end
+
+    context 'when ci: uses different database' do
+      context 'and does not specify database_tasks which indicates using dedicated database' do
+        let(:test_config) do
+          {
+            main: main_database_config,
+            ci: additional_database_config
+          }
+        end
+
+        it_behaves_like 'validates successfully'
+      end
+
+      context 'and does specify database_tasks=false which indicates sharing with main:' do
+        let(:test_config) do
+          {
+            main: main_database_config,
+            ci: additional_database_config.merge(database_tasks: false)
+          }
+        end
+
+        it_behaves_like 'raises an error', /The 'ci' since it is using 'database_tasks: false' should share database with 'main:'/
+      end
+    end
+  end
+
+  %w[db:migrate db:schema:load db:schema:dump].each do |task|
+    context "when running #{task}" do
+      it "does run gitlab:db:validate_config before" do
+        expect(Rake::Task['gitlab:db:validate_config']).to receive(:execute).and_return(true)
+        expect(Rake::Task[task]).to receive(:execute).and_return(true)
+
+        Rake::Task['gitlab:db:validate_config'].reenable
+        run_rake_task(task)
+      end
+    end
+  end
+
+  def with_db_configs(test: test_config)
+    current_configurations = ActiveRecord::Base.configurations # rubocop:disable Database/MultipleDatabases
+    ActiveRecord::Base.configurations = { test: test_config }
+    yield
+  ensure
+    ActiveRecord::Base.configurations = current_configurations
+  end
+end
diff --git a/spec/tasks/gitlab/db_rake_spec.rb b/spec/tasks/gitlab/db_rake_spec.rb
index 8d3ec7b1ee2..73f3b55e12e 100644
--- a/spec/tasks/gitlab/db_rake_spec.rb
+++ b/spec/tasks/gitlab/db_rake_spec.rb
@@ -20,14 +20,6 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
     allow(Rake::Task['db:seed_fu']).to receive(:invoke).and_return(true)
   end
 
-  describe 'clear_all_connections' do
-    it 'calls clear_all_connections!' do
-      expect(ActiveRecord::Base).to receive(:clear_all_connections!)
-
-      run_rake_task('gitlab:db:clear_all_connections')
-    end
-  end
-
   describe 'mark_migration_complete' do
     context 'with a single database' do
       let(:main_model) { ActiveRecord::Base }
@@ -51,7 +43,7 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
       let(:base_models) { { 'main' => main_model, 'ci' => ci_model } }
 
       before do
-        skip_if_multiple_databases_not_setup
+        skip_unless_ci_uses_database_tasks
 
         allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
       end
@@ -80,6 +72,17 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
           run_rake_task('gitlab:db:mark_migration_complete:main', '[123]')
         end
       end
+
+      context 'with geo configured' do
+        before do
+          skip_unless_geo_configured
+        end
+
+        it 'does not create a task for the geo database' do
+          expect { run_rake_task('gitlab:db:mark_migration_complete:geo') }
+            .to raise_error(/Don't know how to build task 'gitlab:db:mark_migration_complete:geo'/)
+        end
+      end
     end
 
     context 'when the migration is already marked complete' do
@@ -122,79 +125,228 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
   end
 
   describe 'configure' do
-    it 'invokes db:migrate when schema has already been loaded' do
-      allow(ActiveRecord::Base.connection).to receive(:tables).and_return(%w[table1 table2])
-      expect(Rake::Task['db:migrate']).to receive(:invoke)
-      expect(Rake::Task['db:structure:load']).not_to receive(:invoke)
-      expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
-      expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
-    end
+    context 'with a single database' do
+      let(:connection) { Gitlab::Database.database_base_models[:main].connection }
+      let(:main_config) { double(:config, name: 'main') }
 
-    it 'invokes db:shema:load and db:seed_fu when schema is not loaded' do
-      allow(ActiveRecord::Base.connection).to receive(:tables).and_return([])
-      expect(Rake::Task['db:structure:load']).to receive(:invoke)
-      expect(Rake::Task['db:seed_fu']).to receive(:invoke)
-      expect(Rake::Task['db:migrate']).not_to receive(:invoke)
-      expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
-    end
+      before do
+        skip_if_multiple_databases_are_setup
+      end
 
-    it 'invokes db:shema:load and db:seed_fu when there is only a single table present' do
-      allow(ActiveRecord::Base.connection).to receive(:tables).and_return(['default'])
-      expect(Rake::Task['db:structure:load']).to receive(:invoke)
-      expect(Rake::Task['db:seed_fu']).to receive(:invoke)
-      expect(Rake::Task['db:migrate']).not_to receive(:invoke)
-      expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
-    end
+      context 'when geo is not configured' do
+        before do
+          allow(ActiveRecord::Base).to receive_message_chain('configurations.configs_for').and_return([main_config])
+        end
 
-    it 'does not invoke any other rake tasks during an error' do
-      allow(ActiveRecord::Base).to receive(:connection).and_raise(RuntimeError, 'error')
-      expect(Rake::Task['db:migrate']).not_to receive(:invoke)
-      expect(Rake::Task['db:structure:load']).not_to receive(:invoke)
-      expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
-      expect { run_rake_task('gitlab:db:configure') }.to raise_error(RuntimeError, 'error')
-      # unstub connection so that the database cleaner still works
-      allow(ActiveRecord::Base).to receive(:connection).and_call_original
-    end
+        context 'when the schema is already loaded' do
+          it 'migrates the database' do
+            allow(connection).to receive(:tables).and_return(%w[table1 table2])
+
+            expect(Rake::Task['db:migrate']).to receive(:invoke)
+            expect(Rake::Task['db:schema:load']).not_to receive(:invoke)
+            expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+
+        context 'when the schema is not loaded' do
+          it 'loads the schema and seeds the database' do
+            allow(connection).to receive(:tables).and_return([])
+
+            expect(Rake::Task['db:schema:load']).to receive(:invoke)
+            expect(Rake::Task['db:seed_fu']).to receive(:invoke)
+            expect(Rake::Task['db:migrate']).not_to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+
+        context 'when only a single table is present' do
+          it 'loads the schema and seeds the database' do
+            allow(connection).to receive(:tables).and_return(['default'])
+
+            expect(Rake::Task['db:schema:load']).to receive(:invoke)
+            expect(Rake::Task['db:seed_fu']).to receive(:invoke)
+            expect(Rake::Task['db:migrate']).not_to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+
+        context 'when loading the schema fails' do
+          it 'does not seed the database' do
+            allow(connection).to receive(:tables).and_return([])
+
+            expect(Rake::Task['db:schema:load']).to receive(:invoke).and_raise('error')
+            expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
+            expect(Rake::Task['db:migrate']).not_to receive(:invoke)
+
+            expect { run_rake_task('gitlab:db:configure') }.to raise_error(RuntimeError, 'error')
+          end
+        end
+
+        context 'SKIP_POST_DEPLOYMENT_MIGRATIONS environment variable set' do
+          let(:rails_paths) { { 'db' => ['db'], 'db/migrate' => ['db/migrate'] } }
+
+          before do
+            stub_env('SKIP_POST_DEPLOYMENT_MIGRATIONS', true)
+
+            # Our environment has already been loaded, so we need to pretend like post_migrations were not
+            allow(Rails.application.config).to receive(:paths).and_return(rails_paths)
+            allow(ActiveRecord::Migrator).to receive(:migrations_paths).and_return(rails_paths['db/migrate'].dup)
+          end
+
+          context 'when the schema is not loaded' do
+            it 'adds the post deployment migration path before schema load' do
+              allow(connection).to receive(:tables).and_return([])
+
+              expect(Gitlab::Database).to receive(:add_post_migrate_path_to_rails).and_call_original
+              expect(Rake::Task['db:schema:load']).to receive(:invoke)
+              expect(Rake::Task['db:seed_fu']).to receive(:invoke)
+              expect(Rake::Task['db:migrate']).not_to receive(:invoke)
+
+              run_rake_task('gitlab:db:configure')
+
+              expect(rails_paths['db/migrate'].include?(File.join(Rails.root, 'db', 'post_migrate'))).to be(true)
+            end
+          end
+
+          context 'when the schema is loaded' do
+            it 'ignores post deployment migrations' do
+              allow(connection).to receive(:tables).and_return(%w[table1 table2])
+
+              expect(Rake::Task['db:migrate']).to receive(:invoke)
+              expect(Gitlab::Database).not_to receive(:add_post_migrate_path_to_rails)
+              expect(Rake::Task['db:schema:load']).not_to receive(:invoke)
+              expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
 
-    it 'does not invoke seed after a failed schema_load' do
-      allow(ActiveRecord::Base.connection).to receive(:tables).and_return([])
-      allow(Rake::Task['db:structure:load']).to receive(:invoke).and_raise(RuntimeError, 'error')
-      expect(Rake::Task['db:structure:load']).to receive(:invoke)
-      expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
-      expect(Rake::Task['db:migrate']).not_to receive(:invoke)
-      expect { run_rake_task('gitlab:db:configure') }.to raise_error(RuntimeError, 'error')
+              run_rake_task('gitlab:db:configure')
+
+              expect(rails_paths['db/migrate'].include?(File.join(Rails.root, 'db', 'post_migrate'))).to be(false)
+            end
+          end
+        end
+      end
+
+      context 'when geo is configured' do
+        context 'when the main database is also configured' do
+          before do
+            skip_unless_geo_configured
+          end
+
+          it 'only configures the main database' do
+            allow(connection).to receive(:tables).and_return(%w[table1 table2])
+
+            expect(Rake::Task['db:migrate:main']).to receive(:invoke)
+
+            expect(Rake::Task['db:migrate:geo']).not_to receive(:invoke)
+            expect(Rake::Task['db:schema:load:geo']).not_to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+      end
     end
 
-    context 'SKIP_POST_DEPLOYMENT_MIGRATIONS environment variable set' do
-      let(:rails_paths) { { 'db' => ['db'], 'db/migrate' => ['db/migrate'] } }
+    context 'with multiple databases' do
+      let(:main_model) { double(:model, connection: double(:connection)) }
+      let(:ci_model) { double(:model, connection: double(:connection)) }
+      let(:base_models) { { 'main' => main_model, 'ci' => ci_model }.with_indifferent_access }
+
+      let(:main_config) { double(:config, name: 'main') }
+      let(:ci_config) { double(:config, name: 'ci') }
 
       before do
-        allow(ENV).to receive(:[]).and_call_original
-        allow(ENV).to receive(:[]).with('SKIP_POST_DEPLOYMENT_MIGRATIONS').and_return true
+        skip_unless_ci_uses_database_tasks
 
-        # Our environment has already been loaded, so we need to pretend like post_migrations were not
-        allow(Rails.application.config).to receive(:paths).and_return(rails_paths)
-        allow(ActiveRecord::Migrator).to receive(:migrations_paths).and_return(rails_paths['db/migrate'].dup)
+        allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
       end
 
-      it 'adds post deployment migrations before schema load if the schema is not already loaded' do
-        allow(ActiveRecord::Base.connection).to receive(:tables).and_return([])
-        expect(Gitlab::Database).to receive(:add_post_migrate_path_to_rails).and_call_original
-        expect(Rake::Task['db:structure:load']).to receive(:invoke)
-        expect(Rake::Task['db:seed_fu']).to receive(:invoke)
-        expect(Rake::Task['db:migrate']).not_to receive(:invoke)
-        expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
-        expect(rails_paths['db/migrate'].include?(File.join(Rails.root, 'db', 'post_migrate'))).to be(true)
+      context 'when geo is not configured' do
+        before do
+          allow(ActiveRecord::Base).to receive_message_chain('configurations.configs_for')
+            .and_return([main_config, ci_config])
+        end
+
+        context 'when no database has the schema loaded' do
+          before do
+            allow(main_model.connection).to receive(:tables).and_return(%w[schema_migrations])
+            allow(ci_model.connection).to receive(:tables).and_return([])
+          end
+
+          it 'loads the schema and seeds all the databases' do
+            expect(Rake::Task['db:schema:load:main']).to receive(:invoke)
+            expect(Rake::Task['db:schema:load:ci']).to receive(:invoke)
+
+            expect(Rake::Task['db:migrate:main']).not_to receive(:invoke)
+            expect(Rake::Task['db:migrate:ci']).not_to receive(:invoke)
+
+            expect(Rake::Task['db:seed_fu']).to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+
+        context 'when both databases have the schema loaded' do
+          before do
+            allow(main_model.connection).to receive(:tables).and_return(%w[table1 table2])
+            allow(ci_model.connection).to receive(:tables).and_return(%w[table1 table2])
+          end
+
+          it 'migrates the databases without seeding them' do
+            expect(Rake::Task['db:migrate:main']).to receive(:invoke)
+            expect(Rake::Task['db:migrate:ci']).to receive(:invoke)
+
+            expect(Rake::Task['db:schema:load:main']).not_to receive(:invoke)
+            expect(Rake::Task['db:schema:load:ci']).not_to receive(:invoke)
+
+            expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+
+        context 'when only one database has the schema loaded' do
+          before do
+            allow(main_model.connection).to receive(:tables).and_return(%w[table1 table2])
+            allow(ci_model.connection).to receive(:tables).and_return([])
+          end
+
+          it 'migrates and loads the schema correctly, without seeding the databases' do
+            expect(Rake::Task['db:migrate:main']).to receive(:invoke)
+            expect(Rake::Task['db:schema:load:main']).not_to receive(:invoke)
+
+            expect(Rake::Task['db:schema:load:ci']).to receive(:invoke)
+            expect(Rake::Task['db:migrate:ci']).not_to receive(:invoke)
+
+            expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
       end
 
-      it 'ignores post deployment migrations  when schema has already been loaded' do
-        allow(ActiveRecord::Base.connection).to receive(:tables).and_return(%w[table1 table2])
-        expect(Rake::Task['db:migrate']).to receive(:invoke)
-        expect(Gitlab::Database).not_to receive(:add_post_migrate_path_to_rails)
-        expect(Rake::Task['db:structure:load']).not_to receive(:invoke)
-        expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
-        expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
-        expect(rails_paths['db/migrate'].include?(File.join(Rails.root, 'db', 'post_migrate'))).to be(false)
+      context 'when geo is configured' do
+        let(:geo_config) { double(:config, name: 'geo') }
+
+        before do
+          skip_unless_geo_configured
+
+          allow(main_model.connection).to receive(:tables).and_return(%w[schema_migrations])
+          allow(ci_model.connection).to receive(:tables).and_return(%w[schema_migrations])
+        end
+
+        it 'does not run tasks against geo' do
+          expect(Rake::Task['db:schema:load:main']).to receive(:invoke)
+          expect(Rake::Task['db:schema:load:ci']).to receive(:invoke)
+          expect(Rake::Task['db:seed_fu']).to receive(:invoke)
+
+          expect(Rake::Task['db:migrate:geo']).not_to receive(:invoke)
+          expect(Rake::Task['db:schema:load:geo']).not_to receive(:invoke)
+
+          run_rake_task('gitlab:db:configure')
+        end
       end
     end
   end
@@ -290,7 +442,7 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
       let(:base_models) { { 'main' => main_model, 'ci' => ci_model } }
 
       before do
-        skip_if_multiple_databases_not_setup
+        skip_unless_ci_uses_database_tasks
 
         allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
 
@@ -319,6 +471,17 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
           run_rake_task('gitlab:db:drop_tables:main')
         end
       end
+
+      context 'with geo configured' do
+        before do
+          skip_unless_geo_configured
+        end
+
+        it 'does not create a task for the geo database' do
+          expect { run_rake_task('gitlab:db:drop_tables:geo') }
+            .to raise_error(/Don't know how to build task 'gitlab:db:drop_tables:geo'/)
+        end
+      end
     end
 
     def expect_objects_to_be_dropped(connection)
@@ -336,38 +499,119 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
     end
   end
 
-  describe 'reindex' do
-    it 'delegates to Gitlab::Database::Reindexing' do
-      expect(Gitlab::Database::Reindexing).to receive(:invoke)
+  describe 'create_dynamic_partitions' do
+    context 'with a single database' do
+      before do
+        skip_if_multiple_databases_are_setup
+      end
+
+      it 'delegates syncing of partitions without limiting databases' do
+        expect(Gitlab::Database::Partitioning).to receive(:sync_partitions)
+
+        run_rake_task('gitlab:db:create_dynamic_partitions')
+      end
+    end
+
+    context 'with multiple databases' do
+      before do
+        skip_unless_ci_uses_database_tasks
+      end
+
+      context 'when running the multi-database variant' do
+        it 'delegates syncing of partitions without limiting databases' do
+          expect(Gitlab::Database::Partitioning).to receive(:sync_partitions)
 
-      run_rake_task('gitlab:db:reindex')
+          run_rake_task('gitlab:db:create_dynamic_partitions')
+        end
+      end
+
+      context 'when running a single-database variant' do
+        it 'delegates syncing of partitions for the chosen database' do
+          expect(Gitlab::Database::Partitioning).to receive(:sync_partitions).with(only_on: 'main')
+
+          run_rake_task('gitlab:db:create_dynamic_partitions:main')
+        end
+      end
     end
 
-    context 'when reindexing is not enabled' do
-      it 'is a no-op' do
-        expect(Gitlab::Database::Reindexing).to receive(:enabled?).and_return(false)
-        expect(Gitlab::Database::Reindexing).not_to receive(:invoke)
+    context 'with geo configured' do
+      before do
+        skip_unless_geo_configured
+      end
 
-        expect { run_rake_task('gitlab:db:reindex') }.to raise_error(SystemExit)
+      it 'does not create a task for the geo database' do
+        expect { run_rake_task('gitlab:db:create_dynamic_partitions:geo') }
+          .to raise_error(/Don't know how to build task 'gitlab:db:create_dynamic_partitions:geo'/)
       end
     end
   end
 
-  databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml
-  ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |database_name|
-    describe "reindex:#{database_name}" do
+  describe 'reindex' do
+    context 'with a single database' do
+      before do
+        skip_if_multiple_databases_are_setup
+      end
+
       it 'delegates to Gitlab::Database::Reindexing' do
-        expect(Gitlab::Database::Reindexing).to receive(:invoke).with(database_name)
+        expect(Gitlab::Database::Reindexing).to receive(:invoke).with(no_args)
 
-        run_rake_task("gitlab:db:reindex:#{database_name}")
+        run_rake_task('gitlab:db:reindex')
       end
 
       context 'when reindexing is not enabled' do
         it 'is a no-op' do
           expect(Gitlab::Database::Reindexing).to receive(:enabled?).and_return(false)
-          expect(Gitlab::Database::Reindexing).not_to receive(:invoke).with(database_name)
+          expect(Gitlab::Database::Reindexing).not_to receive(:invoke)
 
-          expect { run_rake_task("gitlab:db:reindex:#{database_name}") }.to raise_error(SystemExit)
+          expect { run_rake_task('gitlab:db:reindex') }.to raise_error(SystemExit)
+        end
+      end
+    end
+
+    context 'with multiple databases' do
+      let(:base_models) { { 'main' => double(:model), 'ci' => double(:model) } }
+
+      before do
+        skip_if_multiple_databases_not_setup
+
+        allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+      end
+
+      it 'delegates to Gitlab::Database::Reindexing without a specific database' do
+        expect(Gitlab::Database::Reindexing).to receive(:invoke).with(no_args)
+
+        run_rake_task('gitlab:db:reindex')
+      end
+
+      context 'when the single database task is used' do
+        before do
+          skip_unless_ci_uses_database_tasks
+        end
+
+        it 'delegates to Gitlab::Database::Reindexing with a specific database' do
+          expect(Gitlab::Database::Reindexing).to receive(:invoke).with('ci')
+
+          run_rake_task('gitlab:db:reindex:ci')
+        end
+
+        context 'when reindexing is not enabled' do
+          it 'is a no-op' do
+            expect(Gitlab::Database::Reindexing).to receive(:enabled?).and_return(false)
+            expect(Gitlab::Database::Reindexing).not_to receive(:invoke)
+
+            expect { run_rake_task('gitlab:db:reindex:ci') }.to raise_error(SystemExit)
+          end
+        end
+      end
+
+      context 'with geo configured' do
+        before do
+          skip_unless_geo_configured
+        end
+
+        it 'does not create a task for the geo database' do
+          expect { run_rake_task('gitlab:db:reindex:geo') }
+            .to raise_error(/Don't know how to build task 'gitlab:db:reindex:geo'/)
         end
       end
     end
@@ -439,34 +683,77 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
         subject
       end
     end
+
+    describe '#sample_background_migrations' do
+      it 'delegates to the migration runner with a default sample duration' do
+        expect(::Gitlab::Database::Migrations::Runner).to receive_message_chain(:background_migrations, :run_jobs).with(for_duration: 30.minutes)
+
+        run_rake_task('gitlab:db:migration_testing:sample_background_migrations')
+      end
+
+      it 'delegates to the migration runner with a configured sample duration' do
+        expect(::Gitlab::Database::Migrations::Runner).to receive_message_chain(:background_migrations, :run_jobs).with(for_duration: 100.seconds)
+
+        run_rake_task('gitlab:db:migration_testing:sample_background_migrations', '[100]')
+      end
+    end
   end
 
   describe '#execute_batched_migrations' do
-    subject { run_rake_task('gitlab:db:execute_batched_migrations') }
+    subject(:execute_batched_migrations) { run_rake_task('gitlab:db:execute_batched_migrations') }
 
-    let(:migrations) { create_list(:batched_background_migration, 2) }
-    let(:runner) { instance_double('Gitlab::Database::BackgroundMigration::BatchedMigrationRunner') }
+    let(:connections) do
+      {
+        main: instance_double(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter),
+        ci: instance_double(ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
+      }
+    end
+
+    let(:runners) do
+      {
+        main: instance_double('Gitlab::Database::BackgroundMigration::BatchedMigrationRunner'),
+        ci: instance_double('Gitlab::Database::BackgroundMigration::BatchedMigrationRunner')
+      }
+    end
+
+    let(:migrations) do
+      {
+        main: build_list(:batched_background_migration, 1),
+        ci: build_list(:batched_background_migration, 1)
+      }
+    end
 
     before do
-      allow(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive_message_chain(:active, :queue_order).and_return(migrations)
-      allow(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner).to receive(:new).and_return(runner)
+      each_database = class_double('Gitlab::Database::EachDatabase').as_stubbed_const
+
+      allow(each_database).to receive(:each_database_connection)
+        .and_yield(connections[:main], 'main')
+        .and_yield(connections[:ci], 'ci')
+
+      keys = migrations.keys
+      allow(Gitlab::Database::BackgroundMigration::BatchedMigration)
+        .to receive_message_chain(:with_status, :queue_order) { migrations[keys.shift] }
     end
 
     it 'executes all migrations' do
-      migrations.each do |migration|
-        expect(runner).to receive(:run_entire_migration).with(migration)
+      [:main, :ci].each do |name|
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner).to receive(:new)
+          .with(connection: connections[name])
+          .and_return(runners[name])
+
+        expect(runners[name]).to receive(:run_entire_migration).with(migrations[name].first)
       end
 
-      subject
+      execute_batched_migrations
     end
   end
 
   context 'with multiple databases', :reestablished_active_record_base do
     before do
-      skip_if_multiple_databases_not_setup
+      skip_unless_ci_uses_database_tasks
     end
 
-    describe 'db:structure:dump' do
+    describe 'db:structure:dump against a single database' do
       it 'invokes gitlab:db:clean_structure_sql' do
         expect(Rake::Task['gitlab:db:clean_structure_sql']).to receive(:invoke).twice.and_return(true)
 
@@ -474,7 +761,7 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
       end
     end
 
-    describe 'db:schema:dump' do
+    describe 'db:schema:dump against a single database' do
       it 'invokes gitlab:db:clean_structure_sql' do
         expect(Rake::Task['gitlab:db:clean_structure_sql']).to receive(:invoke).once.and_return(true)
 
@@ -482,26 +769,24 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
       end
     end
 
-    describe 'db:migrate' do
-      it 'invokes gitlab:db:create_dynamic_partitions' do
-        expect(Rake::Task['gitlab:db:create_dynamic_partitions']).to receive(:invoke).once.and_return(true)
+    describe 'db:migrate against a single database' do
+      it 'invokes gitlab:db:create_dynamic_partitions for the same database' do
+        expect(Rake::Task['gitlab:db:create_dynamic_partitions:main']).to receive(:invoke).once.and_return(true)
 
         expect { run_rake_task('db:migrate:main') }.not_to raise_error
       end
     end
 
     describe 'db:migrate:geo' do
-      it 'does not invoke gitlab:db:create_dynamic_partitions' do
-        skip 'Skipping because geo database is not setup' unless geo_configured?
+      before do
+        skip_unless_geo_configured
+      end
 
+      it 'does not invoke gitlab:db:create_dynamic_partitions' do
         expect(Rake::Task['gitlab:db:create_dynamic_partitions']).not_to receive(:invoke)
 
         expect { run_rake_task('db:migrate:geo') }.not_to raise_error
       end
-
-      def geo_configured?
-        !!ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: 'geo')
-      end
     end
   end
 
@@ -559,4 +844,20 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
 
     run_rake_task(test_task_name)
   end
+
+  def skip_unless_ci_uses_database_tasks
+    skip "Skipping because database tasks won't run against the ci database" unless ci_database_tasks?
+  end
+
+  def ci_database_tasks?
+    !!ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: 'ci')&.database_tasks?
+  end
+
+  def skip_unless_geo_configured
+    skip 'Skipping because the geo database is not configured' unless geo_configured?
+  end
+
+  def geo_configured?
+    !!ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: 'geo')
+  end
 end
diff --git a/spec/tasks/gitlab/refresh_project_statistics_build_artifacts_size_rake_spec.rb b/spec/tasks/gitlab/refresh_project_statistics_build_artifacts_size_rake_spec.rb
index e57704d0ebe..3495b535cff 100644
--- a/spec/tasks/gitlab/refresh_project_statistics_build_artifacts_size_rake_spec.rb
+++ b/spec/tasks/gitlab/refresh_project_statistics_build_artifacts_size_rake_spec.rb
@@ -11,37 +11,53 @@ RSpec.describe 'gitlab:refresh_project_statistics_build_artifacts_size rake task
     let_it_be(:project_3) { create(:project) }
 
     let(:string_of_ids) { "#{project_1.id} #{project_2.id} #{project_3.id} 999999" }
+    let(:csv_url) { 'https://www.example.com/foo.csv' }
+    let(:csv_body) do
+      <<~BODY
+        PROJECT_ID
+        #{project_1.id}
+        #{project_2.id}
+        #{project_3.id}
+      BODY
+    end
 
     before do
       Rake.application.rake_require('tasks/gitlab/refresh_project_statistics_build_artifacts_size')
 
       stub_const("BUILD_ARTIFACTS_SIZE_REFRESH_ENQUEUE_BATCH_SIZE", 2)
-    end
 
-    context 'when given a list of space-separated IDs through STDIN' do
-      before do
-        allow($stdin).to receive(:tty?).and_return(false)
-        allow($stdin).to receive(:read).and_return(string_of_ids)
-      end
+      stub_request(:get, csv_url).to_return(status: 200, body: csv_body)
+      allow(Kernel).to receive(:sleep).with(1)
+    end
 
+    context 'when given a list of space-separated IDs through rake argument' do
       it 'enqueues the projects for refresh' do
-        expect { run_rake_task(rake_task) }.to output(/Done/).to_stdout
+        expect { run_rake_task(rake_task, csv_url) }.to output(/Done/).to_stdout
 
         expect(Projects::BuildArtifactsSizeRefresh.all.map(&:project)).to match_array([project_1, project_2, project_3])
       end
-    end
 
-    context 'when given a list of space-separated IDs through rake argument' do
-      it 'enqueues the projects for refresh' do
-        expect { run_rake_task(rake_task, string_of_ids) }.to output(/Done/).to_stdout
+      it 'inserts refreshes in batches with a sleep' do
+        expect(Projects::BuildArtifactsSizeRefresh).to receive(:enqueue_refresh).with([project_1, project_2]).ordered
+        expect(Kernel).to receive(:sleep).with(1)
+        expect(Projects::BuildArtifactsSizeRefresh).to receive(:enqueue_refresh).with([project_3]).ordered
 
-        expect(Projects::BuildArtifactsSizeRefresh.all.map(&:project)).to match_array([project_1, project_2, project_3])
+        run_rake_task(rake_task, csv_url)
       end
     end
 
-    context 'when not given any IDs' do
+    context 'when CSV has invalid header' do
+      let(:csv_body) do
+        <<~BODY
+          projectid
+          #{project_1.id}
+          #{project_2.id}
+          #{project_3.id}
+        BODY
+      end
+
       it 'returns an error message' do
-        expect { run_rake_task(rake_task) }.to output(/Please provide a string of space-separated project IDs/).to_stdout
+        expect { run_rake_task(rake_task, csv_url) }.to output(/Project IDs must be listed in the CSV under the header PROJECT_ID/).to_stdout
       end
     end
   end
diff --git a/spec/tasks/gitlab/setup_rake_spec.rb b/spec/tasks/gitlab/setup_rake_spec.rb
index 6e4d5087517..c31546fc259 100644
--- a/spec/tasks/gitlab/setup_rake_spec.rb
+++ b/spec/tasks/gitlab/setup_rake_spec.rb
@@ -6,6 +6,7 @@ RSpec.describe 'gitlab:setup namespace rake tasks', :silence_stdout do
   before do
     Rake.application.rake_require 'active_record/railties/databases'
     Rake.application.rake_require 'tasks/seed_fu'
+    Rake.application.rake_require 'tasks/dev'
     Rake.application.rake_require 'tasks/gitlab/setup'
   end
 
@@ -22,8 +23,6 @@ RSpec.describe 'gitlab:setup namespace rake tasks', :silence_stdout do
     let(:server_service1) { double(:server_service) }
     let(:server_service2) { double(:server_service) }
 
-    let(:connections) { Gitlab::Database.database_base_models.values.map(&:connection) }
-
     before do
       allow(Gitlab).to receive_message_chain('config.repositories.storages').and_return(storages)
 
@@ -98,18 +97,6 @@ RSpec.describe 'gitlab:setup namespace rake tasks', :silence_stdout do
       end
     end
 
-    context 'when the database is not found when terminating connections' do
-      it 'continues setting up the database', :aggregate_failures do
-        expect_gitaly_connections_to_be_checked
-
-        expect(connections).to all(receive(:execute).and_raise(ActiveRecord::NoDatabaseError))
-
-        expect_database_to_be_setup
-
-        setup_task
-      end
-    end
-
     def expect_gitaly_connections_to_be_checked
       expect(Gitlab::GitalyClient::ServerService).to receive(:new).with('name1').and_return(server_service1)
       expect(server_service1).to receive(:info)
@@ -119,13 +106,11 @@ RSpec.describe 'gitlab:setup namespace rake tasks', :silence_stdout do
     end
 
     def expect_connections_to_be_terminated
-      expect(connections).to all(receive(:execute).with(/SELECT pg_terminate_backend/))
+      expect(Rake::Task['dev:terminate_all_connections']).to receive(:invoke)
     end
 
     def expect_connections_not_to_be_terminated
-      connections.each do |connection|
-        expect(connection).not_to receive(:execute)
-      end
+      expect(Rake::Task['dev:terminate_all_connections']).not_to receive(:invoke)
     end
 
     def expect_database_to_be_setup