Display the amount for Hashed Storage migration/rollback correctly63667-hashed-storage-migration-count-correctly

Rake script considers the specified range, if present, to calculate
the affected amount of projects correctly.

2 files changed, 41 insertions, 49 deletions

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!'