Add latest changes from gitlab-org/gitlab@15-5-stable-eev15.5.0-rc42

5 files changed, 269 insertions, 144 deletions

diff --git a/spec/tasks/gitlab/backup_rake_spec.rb b/spec/tasks/gitlab/backup_rake_spec.rb
index dc112b885ae..dc74f25db87 100644
--- a/spec/tasks/gitlab/backup_rake_spec.rb
+++ b/spec/tasks/gitlab/backup_rake_spec.rb
@@ -4,9 +4,10 @@ require 'rake_helper'
 
 RSpec.describe 'gitlab:app namespace rake task', :delete do
   let(:enable_registry) { true }
-  let(:backup_tasks) { %w{db repo uploads builds artifacts pages lfs terraform_state registry packages} }
+  let(:backup_restore_pid_path) { "#{Rails.application.root}/tmp/backup_restore.pid" }
+  let(:backup_tasks) { %w[db repo uploads builds artifacts pages lfs terraform_state registry packages] }
   let(:backup_types) do
-    %w{main_db repositories uploads builds artifacts pages lfs terraform_state registry packages}.tap do |array|
+    %w[main_db repositories uploads builds artifacts pages lfs terraform_state registry packages].tap do |array|
       array.insert(1, 'ci_db') if Gitlab::Database.has_config?(:ci)
     end
   end
@@ -20,11 +21,19 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
   end
 
   def backup_files
-    %w(backup_information.yml artifacts.tar.gz builds.tar.gz lfs.tar.gz terraform_state.tar.gz pages.tar.gz packages.tar.gz)
+    %w[
+      backup_information.yml
+      artifacts.tar.gz
+      builds.tar.gz
+      lfs.tar.gz
+      terraform_state.tar.gz
+      pages.tar.gz
+      packages.tar.gz
+    ]
   end
 
   def backup_directories
-    %w(db repositories)
+    %w[db repositories]
   end
 
   before(:all) do
@@ -58,11 +67,88 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
     end
   end
 
+  describe 'lock parallel backups' do
+    using RSpec::Parameterized::TableSyntax
+
+    context 'when a process is running' do
+      let(:pid_file) { instance_double(File) }
+
+      it 'exits the new process' do
+        allow(File).to receive(:open).and_call_original
+        allow(File).to receive(:open).with(backup_restore_pid_path, any_args).and_yield(pid_file)
+        allow(pid_file).to receive(:read).and_return('123456')
+        allow(pid_file).to receive(:flock).with(any_args)
+
+        expect { run_rake_task('gitlab:backup:create') }.to raise_error(SystemExit).and output(
+          <<~HEREDOC
+            Backup and restore in progress:
+              There is a backup and restore task in progress. Please, try to run the current task once the previous one ends.
+              If there is no other process running, please remove the PID file manually: rm #{backup_restore_pid_path}
+          HEREDOC
+        ).to_stdout
+      end
+    end
+
+    context 'when no processes are running' do
+      let(:progress) { $stdout }
+      let(:pid_file) { instance_double(File, write: 12345) }
+
+      where(:tasks_name, :rake_task) do
+        %w[main_db ci_db] | 'gitlab:backup:db:restore'
+        'repositories'    | 'gitlab:backup:repo:restore'
+        'builds'          | 'gitlab:backup:builds:restore'
+        'uploads'         | 'gitlab:backup:uploads:restore'
+        'artifacts'       | 'gitlab:backup:artifacts:restore'
+        'pages'           | 'gitlab:backup:pages:restore'
+        'lfs'             | 'gitlab:backup:lfs:restore'
+        'terraform_state' | 'gitlab:backup:terraform_state:restore'
+        'registry'        | 'gitlab:backup:registry:restore'
+        'packages'        | 'gitlab:backup:packages:restore'
+      end
+
+      with_them do
+        before do
+          allow(Kernel).to receive(:system).and_return(true)
+          allow(YAML).to receive(:load_file).and_return({ gitlab_version: Gitlab::VERSION })
+          allow(File).to receive(:delete).with(backup_restore_pid_path).and_return(1)
+          allow(File).to receive(:open).and_call_original
+          allow(File).to receive(:open).with(backup_restore_pid_path, any_args).and_yield(pid_file)
+          allow(pid_file).to receive(:read).and_return('')
+          allow(pid_file).to receive(:flock).with(any_args)
+          allow(pid_file).to receive(:write).with(12345).and_return(true)
+          allow(pid_file).to receive(:flush)
+          allow(progress).to receive(:puts).at_least(:once)
+
+          allow_next_instance_of(::Backup::Manager) do |instance|
+            Array(tasks_name).each do |task|
+              allow(instance).to receive(:run_restore_task).with(task)
+            end
+          end
+        end
+
+        it 'locks the PID file' do
+          expect(pid_file).to receive(:flock).with(File::LOCK_EX)
+          expect(pid_file).to receive(:flock).with(File::LOCK_UN)
+
+          run_rake_task(rake_task)
+        end
+
+        it 'deletes the PID file and logs a message' do
+          expect(File).to receive(:delete).with(backup_restore_pid_path)
+          expect(progress).to receive(:puts).with(/-- Deleting backup and restore lock file/)
+
+          run_rake_task(rake_task)
+        end
+      end
+    end
+  end
+
   describe 'backup_restore' do
-    context 'gitlab version' do
+    context 'with gitlab version' do
       before do
         allow(Dir).to receive(:glob).and_return(['1_gitlab_backup.tar'])
         allow(File).to receive(:exist?).and_return(true)
+        allow(File).to receive(:exist?).with(backup_restore_pid_path).and_return(false)
         allow(Kernel).to receive(:system).and_return(true)
         allow(FileUtils).to receive(:cp_r).and_return(true)
         allow(FileUtils).to receive(:mv).and_return(true)
@@ -72,7 +158,7 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
 
       let(:gitlab_version) { Gitlab::VERSION }
 
-      context 'restore with matching gitlab version' do
+      context 'when restore matches gitlab version' do
         before do
           allow(YAML).to receive(:load_file)
             .and_return({ gitlab_version: gitlab_version })
@@ -124,6 +210,7 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         backup_tar = Dir.glob(File.join(Gitlab.config.backup.path, '*_gitlab_backup.tar')).last
         allow(Dir).to receive(:glob).and_return([backup_tar])
         allow(File).to receive(:exist?).and_return(true)
+        allow(File).to receive(:exist?).with(backup_restore_pid_path).and_return(false)
         allow(Kernel).to receive(:system).and_return(true)
         allow(FileUtils).to receive(:cp_r).and_return(true)
         allow(FileUtils).to receive(:mv).and_return(true)
@@ -161,74 +248,42 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
 
     let!(:project) { create(:project, :repository) }
 
-    describe 'backup creation and deletion using custom_hooks' do
-      let(:user_backup_path) { "repositories/#{project.disk_path}" }
-
+    context 'with specific backup tasks' do
       before do
         stub_env('SKIP', 'db')
-        path = Gitlab::GitalyClient::StorageSettings.allow_disk_access do
-          File.join(project.repository.path_to_repo, 'custom_hooks')
-        end
-        FileUtils.mkdir_p(path)
-        FileUtils.touch(File.join(path, "dummy.txt"))
       end
 
-      context 'project uses custom_hooks and successfully creates backup' do
-        it 'creates custom_hooks.tar and project bundle' do
-          expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
-
-          tar_contents, exit_status = Gitlab::Popen.popen(%W{tar -tvf #{backup_tar}})
-
-          expect(exit_status).to eq(0)
-          expect(tar_contents).to match(user_backup_path)
-          expect(tar_contents).to match("#{user_backup_path}/.+/001.custom_hooks.tar")
-          expect(tar_contents).to match("#{user_backup_path}/.+/001.bundle")
-        end
-
-        it 'restores files correctly' do
-          expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
-          expect { run_rake_task('gitlab:backup:restore') }.to output.to_stdout_from_any_process
-
-          repo_path = Gitlab::GitalyClient::StorageSettings.allow_disk_access do
-            project.repository.path
-          end
-          expect(Dir.entries(File.join(repo_path, 'custom_hooks'))).to include("dummy.txt")
+      it 'prints a progress message to stdout' do
+        backup_tasks.each do |task|
+          expect { run_rake_task("gitlab:backup:#{task}:create") }.to output(/Dumping /).to_stdout_from_any_process
         end
       end
 
-      context 'specific backup tasks' do
-        it 'prints a progress message to stdout' do
-          backup_tasks.each do |task|
-            expect { run_rake_task("gitlab:backup:#{task}:create") }.to output(/Dumping /).to_stdout_from_any_process
-          end
-        end
-
-        it 'logs the progress to log file' do
-          ci_database_status = Gitlab::Database.has_config?(:ci) ? "[SKIPPED]" : "[DISABLED]"
-          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping main_database ... [SKIPPED]")
-          expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping ci_database ... #{ci_database_status}")
-          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: "Dumping packages ... done")
-
-          backup_tasks.each do |task|
-            run_rake_task("gitlab:backup:#{task}:create")
-          end
+      it 'logs the progress to log file' do
+        ci_database_status = Gitlab::Database.has_config?(:ci) ? "[SKIPPED]" : "[DISABLED]"
+        expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping main_database ... [SKIPPED]")
+        expect(Gitlab::BackupLogger).to receive(:info).with(message: "Dumping ci_database ... #{ci_database_status}")
+        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: "Dumping packages ... done")
+
+        backup_tasks.each do |task|
+          run_rake_task("gitlab:backup:#{task}:create")
         end
       end
     end
@@ -264,18 +319,18 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
       end
     end
 
-    context 'tar creation' do
-      context 'archive file permissions' do
+    context 'with tar creation' do
+      context 'with archive file permissions' do
         it 'sets correct permissions on the tar file' do
           expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
 
-          expect(File.exist?(backup_tar)).to be_truthy
+          expect(File).to exist(backup_tar)
           expect(File::Stat.new(backup_tar).mode.to_s(8)).to eq('100600')
         end
 
         context 'with custom archive_permissions' do
           before do
-            allow(Gitlab.config.backup).to receive(:archive_permissions).and_return(0651)
+            allow(Gitlab.config.backup).to receive(:archive_permissions).and_return(0o651)
           end
 
           it 'uses the custom permissions' do
@@ -290,11 +345,21 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
 
         tar_contents, exit_status = Gitlab::Popen.popen(
-          %W{tar -tvf #{backup_tar} db uploads.tar.gz repositories builds.tar.gz artifacts.tar.gz pages.tar.gz lfs.tar.gz terraform_state.tar.gz registry.tar.gz packages.tar.gz}
+          %W[
+            tar -tvf #{backup_tar}
+            db
+            uploads.tar.gz
+            repositories
+            builds.tar.gz
+            artifacts.tar.gz
+            pages.tar.gz
+            lfs.tar.gz
+            terraform_state.tar.gz
+            registry.tar.gz
+            packages.tar.gz
+          ]
         )
 
-        puts "CONTENT: #{tar_contents}"
-
         expect(exit_status).to eq(0)
         expect(tar_contents).to match('db')
         expect(tar_contents).to match('uploads.tar.gz')
@@ -306,27 +371,31 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         expect(tar_contents).to match('terraform_state.tar.gz')
         expect(tar_contents).to match('registry.tar.gz')
         expect(tar_contents).to match('packages.tar.gz')
-        expect(tar_contents).not_to match(%r{^.{4,9}[rwx].* (database.sql.gz|uploads.tar.gz|repositories|builds.tar.gz|pages.tar.gz|artifacts.tar.gz|registry.tar.gz)/$})
+        expect(tar_contents).not_to match(%r{^.{4,9}[rwx].* (database.sql.gz|uploads.tar.gz|repositories|builds.tar.gz|
+                                                             pages.tar.gz|artifacts.tar.gz|registry.tar.gz)/$})
       end
 
       it 'deletes temp directories' do
         expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
 
         temp_dirs = Dir.glob(
-          File.join(Gitlab.config.backup.path, '{db,repositories,uploads,builds,artifacts,pages,lfs,terraform_state,registry,packages}')
+          File.join(
+            Gitlab.config.backup.path,
+            '{db,repositories,uploads,builds,artifacts,pages,lfs,terraform_state,registry,packages}'
+          )
         )
 
         expect(temp_dirs).to be_empty
       end
 
-      context 'registry disabled' do
+      context 'when registry is disabled' do
         let(:enable_registry) { false }
 
         it 'does not create registry.tar.gz' do
           expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
 
           tar_contents, exit_status = Gitlab::Popen.popen(
-            %W{tar -tvf #{backup_tar}}
+            %W[tar -tvf #{backup_tar}]
           )
 
           expect(exit_status).to eq(0)
@@ -335,7 +404,7 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
       end
     end
 
-    context 'multiple repository storages' do
+    context 'with multiple repository storages' do
       include StubConfiguration
 
       let(:default_storage_name) { 'default' }
@@ -344,10 +413,10 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
       before do
         # We only need a backup of the repositories for this test
         stub_env('SKIP', 'db,uploads,builds,artifacts,lfs,terraform_state,registry')
-        stub_storage_settings( second_storage_name => {
-          'gitaly_address' => Gitlab.config.repositories.storages.default.gitaly_address,
-          'path' => TestEnv::SECOND_STORAGE_PATH
-        })
+        stub_storage_settings(second_storage_name => {
+                                'gitaly_address' => Gitlab.config.repositories.storages.default.gitaly_address,
+                                'path' => TestEnv::SECOND_STORAGE_PATH
+                              })
       end
 
       shared_examples 'includes repositories in all repository storages' do
@@ -368,27 +437,27 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
           expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
 
           tar_contents, exit_status = Gitlab::Popen.popen(
-            %W{tar -tvf #{backup_tar} repositories}
+            %W[tar -tvf #{backup_tar} repositories]
           )
 
           tar_lines = tar_contents.lines.grep(/\.bundle/)
 
           expect(exit_status).to eq(0)
 
-          [
-            "#{project_a.disk_path}/.+/001.bundle",
-            "#{project_a.disk_path}.wiki/.+/001.bundle",
-            "#{project_a.disk_path}.design/.+/001.bundle",
-            "#{project_b.disk_path}/.+/001.bundle",
-            "#{project_snippet_a.disk_path}/.+/001.bundle",
-            "#{project_snippet_b.disk_path}/.+/001.bundle"
+          %W[
+            #{project_a.disk_path}/.+/001.bundle
+            #{project_a.disk_path}.wiki/.+/001.bundle
+            #{project_a.disk_path}.design/.+/001.bundle
+            #{project_b.disk_path}/.+/001.bundle
+            #{project_snippet_a.disk_path}/.+/001.bundle
+            #{project_snippet_b.disk_path}/.+/001.bundle
           ].each do |repo_name|
             expect(tar_lines).to include(a_string_matching(repo_name))
           end
         end
       end
 
-      context 'no concurrency' do
+      context 'with no concurrency' do
         it_behaves_like 'includes repositories in all repository storages'
       end
 
@@ -400,7 +469,7 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         it_behaves_like 'includes repositories in all repository storages'
       end
 
-      context 'REPOSITORIES_STORAGES set' do
+      context 'when REPOSITORIES_STORAGES is set' do
         before do
           stub_env('REPOSITORIES_STORAGES', default_storage_name)
         end
@@ -422,25 +491,25 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
           expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
 
           tar_contents, exit_status = Gitlab::Popen.popen(
-            %W{tar -tvf #{backup_tar} repositories}
+            %W[tar -tvf #{backup_tar} repositories]
           )
 
           tar_lines = tar_contents.lines.grep(/\.bundle/)
 
           expect(exit_status).to eq(0)
 
-          [
-            "#{project_a.disk_path}/.+/001.bundle",
-            "#{project_a.disk_path}.wiki/.+/001.bundle",
-            "#{project_a.disk_path}.design/.+/001.bundle",
-            "#{project_snippet_a.disk_path}/.+/001.bundle"
+          %W[
+            #{project_a.disk_path}/.+/001.bundle
+            #{project_a.disk_path}.wiki/.+/001.bundle
+            #{project_a.disk_path}.design/.+/001.bundle
+            #{project_snippet_a.disk_path}/.+/001.bundle
           ].each do |repo_name|
             expect(tar_lines).to include(a_string_matching(repo_name))
           end
 
-          [
-            "#{project_b.disk_path}/.+/001.bundle",
-            "#{project_snippet_b.disk_path}/.+/001.bundle"
+          %W[
+            #{project_b.disk_path}/.+/001.bundle
+            #{project_snippet_b.disk_path}/.+/001.bundle
           ].each do |repo_name|
             expect(tar_lines).not_to include(a_string_matching(repo_name))
           end
@@ -448,7 +517,7 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
       end
     end
 
-    context 'concurrency settings' do
+    context 'with concurrency settings' do
       before do
         # We only need a backup of the repositories for this test
         stub_env('SKIP', 'db,uploads,builds,artifacts,lfs,terraform_state,registry')
@@ -463,13 +532,18 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         expect(::Backup::Repositories).to receive(:new)
           .with(anything, strategy: anything, storages: [], paths: [])
           .and_call_original
-        expect(::Backup::GitalyBackup).to receive(:new).with(anything, max_parallelism: 5, storage_parallelism: 2, incremental: false).and_call_original
+        expect(::Backup::GitalyBackup).to receive(:new).with(
+          anything,
+          max_parallelism: 5,
+          storage_parallelism: 2,
+          incremental: false
+        ).and_call_original
 
         expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
       end
     end
 
-    context 'CRON env is set' do
+    context 'when CRON env is set' do
       before do
         stub_env('CRON', '1')
       end
@@ -481,7 +555,7 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
   end
   # backup_create task
 
-  describe "Skipping items in a backup" do
+  describe "skipping items in a backup" do
     before do
       stub_env('SKIP', 'an-unknown-type,repositories,uploads,anotherunknowntype')
 
@@ -492,7 +566,19 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
       expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
 
       tar_contents, _exit_status = Gitlab::Popen.popen(
-        %W{tar -tvf #{backup_tar} db uploads.tar.gz repositories builds.tar.gz artifacts.tar.gz pages.tar.gz lfs.tar.gz terraform_state.tar.gz registry.tar.gz packages.tar.gz}
+        %W[
+          tar -tvf #{backup_tar}
+          db
+          uploads.tar.gz
+          repositories
+          builds.tar.gz
+          artifacts.tar.gz
+          pages.tar.gz
+          lfs.tar.gz
+          terraform_state.tar.gz
+          registry.tar.gz
+          packages.tar.gz
+        ]
       )
 
       expect(tar_contents).to match('db/')
@@ -515,7 +601,7 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         .to receive(:invoke).and_return(true)
 
       expect_next_instance_of(::Backup::Manager) do |instance|
-        (backup_types - %w{repositories uploads}).each do |subtask|
+        (backup_types - %w[repositories uploads]).each do |subtask|
           expect(instance).to receive(:run_restore_task).with(subtask).ordered
         end
         expect(instance).not_to receive(:run_restore_task)
diff --git a/spec/tasks/gitlab/db/lock_writes_rake_spec.rb b/spec/tasks/gitlab/db/lock_writes_rake_spec.rb
index d03e15224cb..ebea644bbf0 100644
--- a/spec/tasks/gitlab/db/lock_writes_rake_spec.rb
+++ b/spec/tasks/gitlab/db/lock_writes_rake_spec.rb
@@ -89,6 +89,26 @@ RSpec.describe 'gitlab:db:lock_writes', :silence_stdout, :reestablished_active_r
       end
     end
 
+    context 'when running in dry_run mode' do
+      before do
+        stub_env('DRY_RUN', 'true')
+      end
+
+      it 'allows writes on the main tables on the ci database' do
+        run_rake_task('gitlab:db:lock_writes')
+        expect do
+          ci_connection.execute("delete from projects")
+        end.not_to raise_error
+      end
+
+      it 'allows writes on the ci tables on the main database' do
+        run_rake_task('gitlab:db:lock_writes')
+        expect do
+          main_connection.execute("delete from ci_builds")
+        end.not_to raise_error
+      end
+    end
+
     context 'multiple shared databases' do
       before do
         allow(::Gitlab::Database).to receive(:db_config_share_with).and_return(nil)
diff --git a/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb b/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb
index f9ebb985255..e95c2e241a8 100644
--- a/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb
+++ b/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb
@@ -43,10 +43,6 @@ RSpec.describe 'gitlab:db:truncate_legacy_tables', :silence_stdout, :reestablish
   end
 
   shared_examples 'truncating legacy tables' do
-    before do
-      allow(ENV).to receive(:[]).and_return(nil)
-    end
-
     context 'when tables are not locked for writes' do
       it 'raises an error when trying to truncate the tables' do
         error_message = /is not locked for writes. Run the rake task gitlab:db:lock_writes first/
@@ -97,7 +93,7 @@ RSpec.describe 'gitlab:db:truncate_legacy_tables', :silence_stdout, :reestablish
 
       context 'when running in dry_run mode' do
         before do
-          allow(ENV).to receive(:[]).with("DRY_RUN").and_return("true")
+          stub_env('DRY_RUN', 'true')
         end
 
         it 'does not truncate any tables' do
@@ -115,7 +111,7 @@ RSpec.describe 'gitlab:db:truncate_legacy_tables', :silence_stdout, :reestablish
 
       context 'when passing until_table parameter via environment variable' do
         before do
-          allow(ENV).to receive(:[]).with("UNTIL_TABLE").and_return(legacy_table)
+          stub_env('UNTIL_TABLE', legacy_table)
         end
 
         it 'sends the table name to TablesTruncate' do
diff --git a/spec/tasks/gitlab/db_rake_spec.rb b/spec/tasks/gitlab/db_rake_spec.rb
index 8f8178cde4d..08bec9fda78 100644
--- a/spec/tasks/gitlab/db_rake_spec.rb
+++ b/spec/tasks/gitlab/db_rake_spec.rb
@@ -665,21 +665,15 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
   end
 
   describe '#migrate_with_instrumentation' do
-    describe '#up' do
-      subject { run_rake_task('gitlab:db:migration_testing:up') }
-
-      it 'delegates to the migration runner' do
-        expect(::Gitlab::Database::Migrations::Runner).to receive_message_chain(:up, :run)
+    let(:runner) { instance_double(::Gitlab::Database::Migrations::Runner) }
 
-        subject
-      end
-    end
-
-    describe '#down' do
-      subject { run_rake_task('gitlab:db:migration_testing:down') }
+    describe '#up (legacy mode)' do
+      subject { run_rake_task('gitlab:db:migration_testing:up') }
 
-      it 'delegates to the migration runner' do
-        expect(::Gitlab::Database::Migrations::Runner).to receive_message_chain(:down, :run)
+      it 'delegates to the migration runner in legacy mode' do
+        expect(::Gitlab::Database::Migrations::Runner).to receive(:up).with(database: 'main', legacy_mode: true)
+                                                                      .and_return(runner)
+        expect(runner).to receive(:run)
 
         subject
       end
@@ -699,31 +693,51 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
       end
     end
 
-    describe '#sample_batched_background_migrations' do
-      let(:batched_runner) { instance_double(::Gitlab::Database::Migrations::TestBatchedBackgroundRunner) }
+    where(:db) do
+      Gitlab::Database::DATABASE_NAMES.map(&:to_sym)
+    end
+
+    with_them do
+      describe '#up' do
+        subject { run_rake_task("gitlab:db:migration_testing:up:#{db}") }
 
-      it 'delegates to the migration runner for the main database with a default sample duration' do
-        expect(::Gitlab::Database::Migrations::Runner).to receive(:batched_background_migrations)
-                                                            .with(for_database: 'main').and_return(batched_runner)
-        expect(batched_runner).to receive(:run_jobs).with(for_duration: 30.minutes)
+        it 'delegates to the migration runner' do
+          expect(::Gitlab::Database::Migrations::Runner).to receive(:up).with(database: db).and_return(runner)
+          expect(runner).to receive(:run)
 
-        run_rake_task('gitlab:db:migration_testing:sample_batched_background_migrations')
+          subject
+        end
       end
 
-      it 'delegates to the migration runner for a specified database with a default sample duration' do
-        expect(::Gitlab::Database::Migrations::Runner).to receive(:batched_background_migrations)
-                                                            .with(for_database: 'ci').and_return(batched_runner)
-        expect(batched_runner).to receive(:run_jobs).with(for_duration: 30.minutes)
+      describe '#down' do
+        subject { run_rake_task("gitlab:db:migration_testing:down:#{db}") }
+
+        it 'delegates to the migration runner' do
+          expect(::Gitlab::Database::Migrations::Runner).to receive(:down).with(database: db).and_return(runner)
+          expect(runner).to receive(:run)
 
-        run_rake_task('gitlab:db:migration_testing:sample_batched_background_migrations', '[ci]')
+          subject
+        end
       end
 
-      it 'delegates to the migration runner for a specified database and sample duration' do
-        expect(::Gitlab::Database::Migrations::Runner).to receive(:batched_background_migrations)
-                                                            .with(for_database: 'ci').and_return(batched_runner)
-        expect(batched_runner).to receive(:run_jobs).with(for_duration: 100.seconds)
+      describe '#sample_batched_background_migrations' do
+        let(:batched_runner) { instance_double(::Gitlab::Database::Migrations::TestBatchedBackgroundRunner) }
+
+        it 'delegates to the migration runner for a specified database with a default sample duration' do
+          expect(::Gitlab::Database::Migrations::Runner).to receive(:batched_background_migrations)
+                                                               .with(for_database: db).and_return(batched_runner)
+          expect(batched_runner).to receive(:run_jobs).with(for_duration: 30.minutes)
 
-        run_rake_task('gitlab:db:migration_testing:sample_batched_background_migrations', '[ci, 100]')
+          run_rake_task("gitlab:db:migration_testing:sample_batched_background_migrations:#{db}")
+        end
+
+        it 'delegates to the migration runner for a specified database and sample duration' do
+          expect(::Gitlab::Database::Migrations::Runner).to receive(:batched_background_migrations)
+                                                              .with(for_database: db).and_return(batched_runner)
+          expect(batched_runner).to receive(:run_jobs).with(for_duration: 100.seconds)
+
+          run_rake_task("gitlab:db:migration_testing:sample_batched_background_migrations:#{db}", '[100]')
+        end
       end
     end
   end
diff --git a/spec/tasks/gitlab/usage_data_rake_spec.rb b/spec/tasks/gitlab/usage_data_rake_spec.rb
index f54d06f406f..7ddba4ceb9b 100644
--- a/spec/tasks/gitlab/usage_data_rake_spec.rb
+++ b/spec/tasks/gitlab/usage_data_rake_spec.rb
@@ -69,6 +69,15 @@ RSpec.describe 'gitlab:usage data take tasks', :silence_stdout do
       expect { run_rake_task('gitlab:usage_data:generate_and_send') }.to output(/.*201.*/).to_stdout
     end
 
+    describe 'generate_ci_template_events' do
+      it "generates #{Gitlab::UsageDataCounters::CiTemplateUniqueCounter::KNOWN_EVENTS_FILE_PATH}" do
+        FileUtils.rm_rf(Gitlab::UsageDataCounters::CiTemplateUniqueCounter::KNOWN_EVENTS_FILE_PATH)
+        run_rake_task('gitlab:usage_data:generate_ci_template_events')
+
+        expect(File.exist?(Gitlab::UsageDataCounters::CiTemplateUniqueCounter::KNOWN_EVENTS_FILE_PATH)).to be true
+      end
+    end
+
     private
 
     def stub_response(url: service_ping_payload_url, body:, status: 201)