Merge branch '63667-hashed-storage-migration-count-correctly' into 'master'

Display the amount for Hashed Storage migration/rollback correctly

Closes #63667

See merge request gitlab-org/gitlab-ce!29996

5 files changed, 118 insertions, 57 deletions

diff --git a/changelogs/unreleased/63667-hashed-storage-migration-count-correctly.yml b/changelogs/unreleased/63667-hashed-storage-migration-count-correctly.yml
new file mode 100644
index 00000000000..ead79be2505
--- /dev/null
+++ b/changelogs/unreleased/63667-hashed-storage-migration-count-correctly.yml
@@ -0,0 +1,5 @@
+---
+title: Display the correct amount of projects being migrated/rolled-back to Hashed Storage when specifying ranges
+merge_request: 29996
+author:
+type: fixed
diff --git a/lib/gitlab/hashed_storage/rake_helper.rb b/lib/gitlab/hashed_storage/rake_helper.rb
index 87a31a37e3f..14727b03ce9 100644
--- a/lib/gitlab/hashed_storage/rake_helper.rb
+++ b/lib/gitlab/hashed_storage/rake_helper.rb
@@ -19,8 +19,12 @@ module Gitlab
         ENV['ID_TO']
       end
 
+      def self.using_ranges?
+        !range_from.nil? && !range_to.nil?
+      end
+
       def self.range_single_item?
-        !range_from.nil? && range_from == range_to
+        using_ranges? && range_from == range_to
       end
 
       # rubocop: disable CodeReuse/ActiveRecord
diff --git a/lib/tasks/gitlab/storage.rake b/lib/tasks/gitlab/storage.rake
index 954f827f716..ccc96b7edfb 100644
--- a/lib/tasks/gitlab/storage.rake
+++ b/lib/tasks/gitlab/storage.rake
@@ -3,50 +3,44 @@ namespace :gitlab do
     desc 'GitLab | Storage | Migrate existing projects to Hashed Storage'
     task migrate_to_hashed: :environment do
       if Gitlab::Database.read_only?
-        warn 'This task requires database write access. Exiting.'
-
-        next
+        abort 'This task requires database write access. Exiting.'
       end
 
       storage_migrator = Gitlab::HashedStorage::Migrator.new
       helper = Gitlab::HashedStorage::RakeHelper
 
       if storage_migrator.rollback_pending?
-        warn "There is already a rollback operation in progress, " \
+        abort "There is already a rollback operation in progress, " \
              "running a migration at the same time may have unexpected consequences."
-
-        next
       end
 
       if helper.range_single_item?
         project = Project.with_unmigrated_storage.find_by(id: helper.range_from)
 
         unless project
-          warn "There are no projects requiring storage migration with ID=#{helper.range_from}"
-
-          next
+          abort "There are no projects requiring storage migration with ID=#{helper.range_from}"
         end
 
         puts "Enqueueing storage migration of #{project.full_path} (ID=#{project.id})..."
         storage_migrator.migrate(project)
+      else
+        legacy_projects_count = if helper.using_ranges?
+                                  Project.with_unmigrated_storage.id_in(helper.range_from..helper.range_to).count
+                                else
+                                  Project.with_unmigrated_storage.count
+                                end
+
+        if legacy_projects_count == 0
+          abort 'There are no projects requiring storage migration. Nothing to do!'
+        end
 
-        next
-      end
-
-      legacy_projects_count = Project.with_unmigrated_storage.count
-
-      if legacy_projects_count == 0
-        warn 'There are no projects requiring storage migration. Nothing to do!'
-
-        next
-      end
-
-      print "Enqueuing migration of #{legacy_projects_count} projects in batches of #{helper.batch_size}"
+        print "Enqueuing migration of #{legacy_projects_count} projects in batches of #{helper.batch_size}"
 
-      helper.project_id_batches_migration do |start, finish|
-        storage_migrator.bulk_schedule_migration(start: start, finish: finish)
+        helper.project_id_batches_migration do |start, finish|
+          storage_migrator.bulk_schedule_migration(start: start, finish: finish)
 
-        print '.'
+          print '.'
+        end
       end
 
       puts ' Done!'
@@ -55,50 +49,44 @@ namespace :gitlab do
     desc 'GitLab | Storage | Rollback existing projects to Legacy Storage'
     task rollback_to_legacy: :environment do
       if Gitlab::Database.read_only?
-        warn 'This task requires database write access. Exiting.'
-
-        next
+        abort 'This task requires database write access. Exiting.'
       end
 
       storage_migrator = Gitlab::HashedStorage::Migrator.new
       helper = Gitlab::HashedStorage::RakeHelper
 
       if storage_migrator.migration_pending?
-        warn "There is already a migration operation in progress, " \
+        abort "There is already a migration operation in progress, " \
              "running a rollback at the same time may have unexpected consequences."
-
-        next
       end
 
       if helper.range_single_item?
         project = Project.with_storage_feature(:repository).find_by(id: helper.range_from)
 
         unless project
-          warn "There are no projects that can be rolledback with ID=#{helper.range_from}"
-
-          next
+          abort "There are no projects that can be rolledback with ID=#{helper.range_from}"
         end
 
         puts "Enqueueing storage rollback of #{project.full_path} (ID=#{project.id})..."
         storage_migrator.rollback(project)
+      else
+        hashed_projects_count = if helper.using_ranges?
+                                  Project.with_storage_feature(:repository).id_in(helper.range_from..helper.range_to).count
+                                else
+                                  Project.with_storage_feature(:repository).count
+                                end
+
+        if hashed_projects_count == 0
+          abort 'There are no projects that can have storage rolledback. Nothing to do!'
+        end
 
-        next
-      end
-
-      hashed_projects_count = Project.with_storage_feature(:repository).count
-
-      if hashed_projects_count == 0
-        warn 'There are no projects that can have storage rolledback. Nothing to do!'
-
-        next
-      end
-
-      print "Enqueuing rollback of #{hashed_projects_count} projects in batches of #{helper.batch_size}"
+        print "Enqueuing rollback of #{hashed_projects_count} projects in batches of #{helper.batch_size}"
 
-      helper.project_id_batches_rollback do |start, finish|
-        storage_migrator.bulk_schedule_rollback(start: start, finish: finish)
+        helper.project_id_batches_rollback do |start, finish|
+          storage_migrator.bulk_schedule_rollback(start: start, finish: finish)
 
-        print '.'
+          print '.'
+        end
       end
 
       puts ' Done!'
diff --git a/spec/support/matchers/abort_matcher.rb b/spec/support/matchers/abort_matcher.rb
new file mode 100644
index 00000000000..ce1dd140210
--- /dev/null
+++ b/spec/support/matchers/abort_matcher.rb
@@ -0,0 +1,46 @@
+RSpec::Matchers.define :abort_execution do
+  match do |code_block|
+    @captured_stderr = StringIO.new
+    original_stderr = $stderr
+    $stderr = @captured_stderr
+
+    code_block.call
+
+    false
+  rescue SystemExit => e
+    captured = @captured_stderr.string.chomp
+    @actual_exit_code = e.status
+    break false unless e.status == 1
+
+    if @message
+      if @message.is_a? String
+        @message == captured
+      elsif @message.is_a? Regexp
+        @message.match?(captured)
+      else
+        raise ArgumentError, 'with_message must be either a String or a Regular Expression'
+      end
+    end
+
+  ensure
+    $stderr = original_stderr
+  end
+
+  chain :with_message do |message|
+    @message = message
+  end
+
+  failure_message do |block|
+    unless @actual_exit_code
+      break "expected #{block} to abort with '#{@message}' but didnt call abort."
+    end
+
+    if @actual_exit_code != 1
+      break "expected #{block} to abort with: '#{@message}' but exited with success instead."
+    end
+
+    "expected #{block} to abort with: '#{@message}' \n but received: '#{@captured_stderr.string.chomp}' instead."
+  end
+
+  supports_block_expectations
+end
diff --git a/spec/tasks/gitlab/storage_rake_spec.rb b/spec/tasks/gitlab/storage_rake_spec.rb
index 4b04d9cec39..0e47408fc72 100644
--- a/spec/tasks/gitlab/storage_rake_spec.rb
+++ b/spec/tasks/gitlab/storage_rake_spec.rb
@@ -50,7 +50,7 @@ describe 'rake gitlab:storage:*', :sidekiq do
 
         expect(Project).not_to receive(:with_unmigrated_storage)
 
-        expect { run_rake_task(task) }.to output(/This task requires database write access. Exiting./).to_stderr
+        expect { run_rake_task(task) }.to abort_execution.with_message(/This task requires database write access. Exiting./)
       end
     end
   end
@@ -96,7 +96,7 @@ describe 'rake gitlab:storage:*', :sidekiq do
 
           expect(Project).not_to receive(:with_unmigrated_storage)
 
-          expect { run_rake_task(task) }.to output(/There is already a rollback operation in progress/).to_stderr
+          expect { run_rake_task(task) }.to abort_execution.with_message(/There is already a rollback operation in progress/)
         end
       end
     end
@@ -105,14 +105,23 @@ describe 'rake gitlab:storage:*', :sidekiq do
       it 'does nothing' do
         expect(::HashedStorage::MigratorWorker).not_to receive(:perform_async)
 
-        run_rake_task(task)
+        expect { run_rake_task(task) }.to abort_execution.with_message('There are no projects requiring storage migration. Nothing to do!')
       end
     end
 
     context 'with 3 legacy projects' do
       let(:projects) { create_list(:project, 3, :legacy_storage) }
 
-      it_behaves_like "handles custom BATCH env var", ::HashedStorage::MigratorWorker
+      it 'enqueues migrations and count projects correctly' do
+        projects.map(&:id).sort.tap do |ids|
+          stub_env('ID_FROM', ids[0])
+          stub_env('ID_TO', ids[1])
+        end
+
+        expect { run_rake_task(task) }.to output(/Enqueuing migration of 2 projects in batches/).to_stdout
+      end
+
+      it_behaves_like 'handles custom BATCH env var', ::HashedStorage::MigratorWorker
     end
 
     context 'with same id in range' do
@@ -120,7 +129,7 @@ describe 'rake gitlab:storage:*', :sidekiq do
         stub_env('ID_FROM', 99999)
         stub_env('ID_TO', 99999)
 
-        expect { run_rake_task(task) }.to output(/There are no projects requiring storage migration with ID=99999/).to_stderr
+        expect { run_rake_task(task) }.to abort_execution.with_message(/There are no projects requiring storage migration with ID=99999/)
       end
 
       it 'displays a message when project exists but its already migrated' do
@@ -128,7 +137,7 @@ describe 'rake gitlab:storage:*', :sidekiq do
         stub_env('ID_FROM', project.id)
         stub_env('ID_TO', project.id)
 
-        expect { run_rake_task(task) }.to output(/There are no projects requiring storage migration with ID=#{project.id}/).to_stderr
+        expect { run_rake_task(task) }.to abort_execution.with_message(/There are no projects requiring storage migration with ID=#{project.id}/)
       end
 
       it 'enqueues migration when project can be found' do
@@ -153,7 +162,7 @@ describe 'rake gitlab:storage:*', :sidekiq do
 
           expect(Project).not_to receive(:with_unmigrated_storage)
 
-          expect { run_rake_task(task) }.to output(/There is already a migration operation in progress/).to_stderr
+          expect { run_rake_task(task) }.to abort_execution.with_message(/There is already a migration operation in progress/)
         end
       end
     end
@@ -162,13 +171,22 @@ describe 'rake gitlab:storage:*', :sidekiq do
       it 'does nothing' do
         expect(::HashedStorage::RollbackerWorker).not_to receive(:perform_async)
 
-        run_rake_task(task)
+        expect { run_rake_task(task) }.to abort_execution.with_message('There are no projects that can have storage rolledback. Nothing to do!')
       end
     end
 
     context 'with 3 hashed projects' do
       let(:projects) { create_list(:project, 3) }
 
+      it 'enqueues migrations and count projects correctly' do
+        projects.map(&:id).sort.tap do |ids|
+          stub_env('ID_FROM', ids[0])
+          stub_env('ID_TO', ids[1])
+        end
+
+        expect { run_rake_task(task) }.to output(/Enqueuing rollback of 2 projects in batches/).to_stdout
+      end
+
       it_behaves_like "handles custom BATCH env var", ::HashedStorage::RollbackerWorker
     end
   end