Prepare rake task for storage rollback

We are keeping compatibility with existing scheduled jobs.

6 files changed, 37 insertions, 20 deletions

diff --git a/app/workers/storage_migrator_worker.rb b/app/workers/storage_migrator_worker.rb
index fa76fbac55c..714c4eed5b2 100644
--- a/app/workers/storage_migrator_worker.rb
+++ b/app/workers/storage_migrator_worker.rb
@@ -3,8 +3,14 @@
 class StorageMigratorWorker
   include ApplicationWorker
 
-  def perform(start, finish)
+  # @param [Integer] start initial ID of the batch
+  # @param [Integer] finish last ID of the batch
+  # @param [String] operation the operation to be performed: ['migrate', 'rollback']
+  def perform(start, finish, operation = :migrate)
+    # when scheduling a job symbols are converted to string, we need to convert back
+    operation = operation.to_sym if operation
+
     migrator = Gitlab::HashedStorage::Migrator.new
-    migrator.bulk_migrate(start, finish)
+    migrator.bulk_migrate(start: start, finish: finish, operation: operation)
   end
 end
diff --git a/lib/gitlab/hashed_storage/migrator.rb b/lib/gitlab/hashed_storage/migrator.rb
index 1f29cf10cad..4794b1978db 100644
--- a/lib/gitlab/hashed_storage/migrator.rb
+++ b/lib/gitlab/hashed_storage/migrator.rb
@@ -11,10 +11,11 @@ module Gitlab
 
       # Schedule a range of projects to be bulk migrated with #bulk_migrate asynchronously
       #
-      # @param [Object] start first project id for the range
-      # @param [Object] finish last project id for the range
-      def bulk_schedule(start, finish)
-        StorageMigratorWorker.perform_async(start, finish)
+      # @param [Integer] start first project id for the range
+      # @param [Integer] finish last project id for the range
+      # @param [Symbol] operation [:migrate, :rollback]
+      def bulk_schedule(start:, finish:, operation: :migrate)
+        StorageMigratorWorker.perform_async(start, finish, operation)
       end
 
       # Start migration of projects from specified range
@@ -22,21 +23,27 @@ module Gitlab
       # Flagging a project to be migrated is a synchronous action,
       # but the migration runs through async jobs
       #
-      # @param [Object] start first project id for the range
-      # @param [Object] finish last project id for the range
+      # @param [Integer] start first project id for the range
+      # @param [Integer] finish last project id for the range
+      # @param [Symbol] operation [:migrate, :rollback]
       # rubocop: disable CodeReuse/ActiveRecord
-      def bulk_migrate(start, finish)
+      def bulk_migrate(start:, finish:, operation: :migrate)
         projects = build_relation(start, finish)
 
         projects.with_route.find_each(batch_size: BATCH_SIZE) do |project|
-          migrate(project)
+          case operation
+          when :migrate
+            migrate(project)
+          when :rollback
+            rollback(project)
+          end
         end
       end
       # rubocop: enable CodeReuse/ActiveRecord
 
       # Flag a project to be migrated
       #
-      # @param [Object] project that will be migrated
+      # @param [Project] project that will be migrated
       def migrate(project)
         Rails.logger.info "Starting storage migration of #{project.full_path} (ID=#{project.id})..."
 
@@ -45,6 +52,10 @@ module Gitlab
         Rails.logger.error("#{err.message} migrating storage of #{project.full_path} (ID=#{project.id}), trace - #{err.backtrace}")
       end
 
+      def rollback(project)
+        # TODO: implement rollback strategy
+      end
+
       private
 
       # rubocop: disable CodeReuse/ActiveRecord
diff --git a/lib/tasks/gitlab/storage.rake b/lib/tasks/gitlab/storage.rake
index 09dc3aa9882..8212be90b97 100644
--- a/lib/tasks/gitlab/storage.rake
+++ b/lib/tasks/gitlab/storage.rake
@@ -37,7 +37,7 @@ namespace :gitlab do
       print "Enqueuing migration of #{legacy_projects_count} projects in batches of #{helper.batch_size}"
 
       helper.project_id_batches do |start, finish|
-        storage_migrator.bulk_schedule(start, finish)
+        storage_migrator.bulk_schedule(start: start, finish: finish, operation: :migrate)
 
         print '.'
       end
diff --git a/spec/lib/gitlab/hashed_storage/migrator_spec.rb b/spec/lib/gitlab/hashed_storage/migrator_spec.rb
index 01d43ed00a2..6d68cb0744b 100644
--- a/spec/lib/gitlab/hashed_storage/migrator_spec.rb
+++ b/spec/lib/gitlab/hashed_storage/migrator_spec.rb
@@ -4,7 +4,7 @@ describe Gitlab::HashedStorage::Migrator do
   describe '#bulk_schedule' do
     it 'schedules job to StorageMigratorWorker' do
       Sidekiq::Testing.fake! do
-        expect { subject.bulk_schedule(1, 5) }.to change(StorageMigratorWorker.jobs, :size).by(1)
+        expect { subject.bulk_schedule(start: 1, finish: 5) }.to change(StorageMigratorWorker.jobs, :size).by(1)
       end
     end
   end
@@ -15,13 +15,13 @@ describe Gitlab::HashedStorage::Migrator do
 
     it 'enqueue jobs to ProjectMigrateHashedStorageWorker' do
       Sidekiq::Testing.fake! do
-        expect { subject.bulk_migrate(ids.min, ids.max) }.to change(ProjectMigrateHashedStorageWorker.jobs, :size).by(2)
+        expect { subject.bulk_migrate(start: ids.min, finish: ids.max) }.to change(ProjectMigrateHashedStorageWorker.jobs, :size).by(2)
       end
     end
 
     it 'rescues and log exceptions' do
       allow_any_instance_of(Project).to receive(:migrate_to_hashed_storage!).and_raise(StandardError)
-      expect { subject.bulk_migrate(ids.min, ids.max) }.not_to raise_error
+      expect { subject.bulk_migrate(start: ids.min, finish: ids.max) }.not_to raise_error
     end
 
     it 'delegates each project in specified range to #migrate' do
@@ -29,12 +29,12 @@ describe Gitlab::HashedStorage::Migrator do
         expect(subject).to receive(:migrate).with(project)
       end
 
-      subject.bulk_migrate(ids.min, ids.max)
+      subject.bulk_migrate(start: ids.min, finish: ids.max)
     end
 
     it 'has migrated projects set as writable' do
       perform_enqueued_jobs do
-        subject.bulk_migrate(ids.min, ids.max)
+        subject.bulk_migrate(start: ids.min, finish: ids.max)
       end
 
       projects.each do |project|
diff --git a/spec/tasks/gitlab/storage_rake_spec.rb b/spec/tasks/gitlab/storage_rake_spec.rb
index be902d7c679..3b8d5ad7015 100644
--- a/spec/tasks/gitlab/storage_rake_spec.rb
+++ b/spec/tasks/gitlab/storage_rake_spec.rb
@@ -74,7 +74,7 @@ describe 'rake gitlab:storage:*' do
 
         it 'enqueues one StorageMigratorWorker per project' do
           projects.each do |project|
-            expect(StorageMigratorWorker).to receive(:perform_async).with(project.id, project.id)
+            expect(StorageMigratorWorker).to receive(:perform_async).with(project.id, project.id, :migrate)
           end
 
           run_rake_task(task)
@@ -89,7 +89,7 @@ describe 'rake gitlab:storage:*' do
         it 'enqueues one StorageMigratorWorker per 2 projects' do
           projects.map(&:id).sort.each_slice(2) do |first, last|
             last ||= first
-            expect(StorageMigratorWorker).to receive(:perform_async).with(first, last)
+            expect(StorageMigratorWorker).to receive(:perform_async).with(first, last, :migrate)
           end
 
           run_rake_task(task)
diff --git a/spec/workers/storage_migrator_worker_spec.rb b/spec/workers/storage_migrator_worker_spec.rb
index 808084c8f7c..63f3c1d1c59 100644
--- a/spec/workers/storage_migrator_worker_spec.rb
+++ b/spec/workers/storage_migrator_worker_spec.rb
@@ -7,7 +7,7 @@ describe StorageMigratorWorker do
 
   describe '#perform' do
     it 'delegates to MigratorService' do
-      expect_any_instance_of(Gitlab::HashedStorage::Migrator).to receive(:bulk_migrate).with(5, 10)
+      expect_any_instance_of(Gitlab::HashedStorage::Migrator).to receive(:bulk_migrate).with(start: 5, finish: 10, operation: :migrate)
 
       worker.perform(5, 10)
     end