Fix the rake task to actually run and schedule the jobs for rollback53966-make-hashed-storage-migration-safer-and-more-inviting

2 files changed, 14 insertions, 3 deletions

diff --git a/lib/gitlab/hashed_storage/rake_helper.rb b/lib/gitlab/hashed_storage/rake_helper.rb
index 38f552fab03..87a31a37e3f 100644
--- a/lib/gitlab/hashed_storage/rake_helper.rb
+++ b/lib/gitlab/hashed_storage/rake_helper.rb
@@ -24,7 +24,7 @@ module Gitlab
       end
 
       # rubocop: disable CodeReuse/ActiveRecord
-      def self.project_id_batches(&block)
+      def self.project_id_batches_migration(&block)
         Project.with_unmigrated_storage.in_batches(of: batch_size, start: range_from, finish: range_to) do |relation| # rubocop: disable Cop/InBatches
           ids = relation.pluck(:id)
 
@@ -34,6 +34,16 @@ module Gitlab
       # rubocop: enable CodeReuse/ActiveRecord
 
       # rubocop: disable CodeReuse/ActiveRecord
+      def self.project_id_batches_rollback(&block)
+        Project.with_storage_feature(:repository).in_batches(of: batch_size, start: range_from, finish: range_to) do |relation| # rubocop: disable Cop/InBatches
+          ids = relation.pluck(:id)
+
+          yield ids.min, ids.max
+        end
+      end
+      # rubocop: enable CodeReuse/ActiveRecord
+
+      # rubocop: disable CodeReuse/ActiveRecord
       def self.legacy_attachments_relation
         Upload.joins(<<~SQL).where('projects.storage_version < :version OR projects.storage_version IS NULL', version: Project::HASHED_STORAGE_FEATURES[:attachments])
           JOIN projects
diff --git a/lib/tasks/gitlab/storage.rake b/lib/tasks/gitlab/storage.rake
index 1e64a481151..a2136ce1b92 100644
--- a/lib/tasks/gitlab/storage.rake
+++ b/lib/tasks/gitlab/storage.rake
@@ -36,7 +36,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|
+      helper.project_id_batches_migration do |start, finish|
         storage_migrator.bulk_schedule_migration(start: start, finish: finish)
 
         print '.'
@@ -81,7 +81,8 @@ namespace :gitlab do
 
       print "Enqueuing rollback of #{hashed_projects_count} projects in batches of #{helper.batch_size}"
 
-      helper.project_id_batches do |start, finish|
+      helper.project_id_batches_rollback do |start, finish|
+        puts "Start: #{start} FINISH: #{finish}"
         storage_migrator.bulk_schedule_rollback(start: start, finish: finish)
 
         print '.'