Add rake task for cleaning migrated pages deployments

it can be used if we need to restart migration from scratch

app/models/pages_deployment.rb

@@ -2,6 +2,7 @@
 
 # PagesDeployment stores a zip archive containing GitLab Pages web-site
 class PagesDeployment < ApplicationRecord
+  include EachBatch
   include FileStoreMounter
 
   MIGRATED_FILE_NAME = "_migrated.zip"

lib/tasks/gitlab/pages.rake

@@ -6,7 +6,6 @@ namespace :gitlab do
   namespace :pages do
     desc "GitLab | Pages | Migrate legacy storage to zip format"
     task migrate_legacy_storage: :gitlab_environment do
-      logger = Logger.new(STDOUT)
       logger.info('Starting to migrate legacy pages storage to zip deployments')
 
       result = ::Pages::MigrateFromLegacyStorageService.new(logger, migration_threads, batch_size).execute
@@ -16,6 +15,26 @@ namespace :gitlab do
       logger.info("- The #{result[:errored]} projects failed to be migrated")
     end
 
+    desc "GitLab | Pages | DANGER: Removes data which was migrated from legacy storage on zip storage. Can be used if some bugs in migration are discovered and migration needs to be restarted from scratch."
+    task clean_migrated_zip_storage: :gitlab_environment do
+      destroyed_deployments = 0
+
+      logger.info("Starting to delete migrated pages deployments")
+
+      ::PagesDeployment.migrated_from_legacy_storage.each_batch(of: batch_size) do |batch|
+        destroyed_deployments += batch.count
+
+        # we need to destroy associated files, so can't use delete_all
+        batch.destroy_all # rubocop: disable Cop/DestroyAll
+
+        logger.info("#{destroyed_deployments} deployments were deleted")
+      end
+    end
+
+    def logger
+      @logger ||= Logger.new(STDOUT)
+    end
+
     def migration_threads
       ENV.fetch('PAGES_MIGRATION_THREADS', '3').to_i
     end

spec/factories/pages_deployments.rb

@@ -4,12 +4,20 @@ FactoryBot.define do
   factory :pages_deployment, class: 'PagesDeployment' do
     project
 
-    after(:build) do |deployment, _evaluator|
-      filepath = Rails.root.join("spec/fixtures/pages.zip")
+    transient do
+      filename { nil }
+    end
+
+    trait(:migrated) do
+      filename { PagesDeployment::MIGRATED_FILE_NAME }
+    end
+
+    after(:build) do |deployment, evaluator|
+      file = UploadedFile.new("spec/fixtures/pages.zip", filename: evaluator.filename)
 
-      deployment.file = fixture_file_upload(filepath)
-      deployment.file_sha256 = Digest::SHA256.file(filepath).hexdigest
-      ::Zip::File.open(filepath) do |zip_archive|
+      deployment.file = file
+      deployment.file_sha256 = Digest::SHA256.file(file.path).hexdigest
+      ::Zip::File.open(file.path) do |zip_archive|
         deployment.file_count = zip_archive.count
       end
     end

spec/tasks/gitlab/pages_rake_spec.rb

@@ -2,38 +2,58 @@
 
 require 'rake_helper'
 
-RSpec.describe 'gitlab:pages:migrate_legacy_storagerake task' do
+RSpec.describe 'gitlab:pages' do
   before(:context) do
     Rake.application.rake_require 'tasks/gitlab/pages'
   end
 
-  subject { run_rake_task('gitlab:pages:migrate_legacy_storage') }
+  describe 'migrate_legacy_storage task' do
+    subject { run_rake_task('gitlab:pages:migrate_legacy_storage') }
 
-  it 'calls migration service' do
-    expect_next_instance_of(::Pages::MigrateFromLegacyStorageService, anything, 3, 10) do |service|
-      expect(service).to receive(:execute).and_call_original
+    it 'calls migration service' do
+      expect_next_instance_of(::Pages::MigrateFromLegacyStorageService, anything, 3, 10) do |service|
+        expect(service).to receive(:execute).and_call_original
+      end
+
+      subject
     end
 
-    subject
-  end
+    it 'uses PAGES_MIGRATION_THREADS environment variable' do
+      stub_env('PAGES_MIGRATION_THREADS', '5')
 
-  it 'uses PAGES_MIGRATION_THREADS environment variable' do
-    stub_env('PAGES_MIGRATION_THREADS', '5')
+      expect_next_instance_of(::Pages::MigrateFromLegacyStorageService, anything, 5, 10) do |service|
+        expect(service).to receive(:execute).and_call_original
+      end
 
-    expect_next_instance_of(::Pages::MigrateFromLegacyStorageService, anything, 5, 10) do |service|
-      expect(service).to receive(:execute).and_call_original
+      subject
     end
 
-    subject
-  end
+    it 'uses PAGES_MIGRATION_BATCH_SIZE environment variable' do
+      stub_env('PAGES_MIGRATION_BATCH_SIZE', '100')
 
-  it 'uses PAGES_MIGRATION_BATCH_SIZE environment variable' do
-    stub_env('PAGES_MIGRATION_BATCH_SIZE', '100')
+      expect_next_instance_of(::Pages::MigrateFromLegacyStorageService, anything, 3, 100) do |service|
+        expect(service).to receive(:execute).and_call_original
+      end
 
-    expect_next_instance_of(::Pages::MigrateFromLegacyStorageService, anything, 3, 100) do |service|
-      expect(service).to receive(:execute).and_call_original
+      subject
     end
+  end
+
+  describe 'clean_migrated_zip_storage task' do
+    it 'removes only migrated deployments' do
+      regular_deployment = create(:pages_deployment)
+      migrated_deployment = create(:pages_deployment, :migrated)
 
-    subject
+      regular_deployment.project.update_pages_deployment!(regular_deployment)
+      migrated_deployment.project.update_pages_deployment!(migrated_deployment)
+
+      expect(PagesDeployment.all).to contain_exactly(regular_deployment, migrated_deployment)
+
+      run_rake_task('gitlab:pages:clean_migrated_zip_storage')
+
+      expect(PagesDeployment.all).to contain_exactly(regular_deployment)
+      expect(PagesDeployment.find_by_id(regular_deployment.id)).not_to be_nil
+      expect(PagesDeployment.find_by_id(migrated_deployment.id)).to be_nil
+    end
   end
 end