Adds a button to retry a failed migration

Allows admins to retry batched background migrations

Related to https://gitlab.com/gitlab-org/gitlab/-/issues/326760

Changelog: added

app/controllers/admin/background_migrations_controller.rb

@@ -29,9 +29,16 @@ class Admin::BackgroundMigrationsController < Admin::ApplicationController
     redirect_back fallback_location: { action: 'index' }
   end
 
+  def retry
+    migration = batched_migration_class.find(params[:id])
+    migration.retry_failed_jobs! if migration.failed?
+
+    redirect_back fallback_location: { action: 'index' }
+  end
+
   private
 
   def batched_migration_class
-    Gitlab::Database::BackgroundMigration::BatchedMigration
+    @batched_migration_class ||= Gitlab::Database::BackgroundMigration::BatchedMigration
   end
 end

app/views/admin/background_migrations/_migration.html.haml

@@ -17,3 +17,7 @@
       = button_to resume_admin_background_migration_path(migration),
         class: 'gl-button btn btn-icon has-tooltip', title: _('Resume'), 'aria-label' => _('Resume') do
         = sprite_icon('play', css_class: 'gl-button-icon gl-icon')
+    - elsif migration.failed?
+      = button_to retry_admin_background_migration_path(migration),
+        class: 'gl-button btn btn-icon has-tooltip', title: _('Retry'), 'aria-label' => _('Retry') do
+        = sprite_icon('retry', css_class: 'gl-button-icon gl-icon')

config/routes/admin.rb

@@ -93,6 +93,7 @@ namespace :admin do
     member do
       post :pause
       post :resume
+      post :retry
     end
   end
 

lib/gitlab/database/background_migration/batched_job.rb

@@ -4,6 +4,7 @@ module Gitlab
   module Database
     module BackgroundMigration
       class BatchedJob < ActiveRecord::Base # rubocop:disable Rails/ApplicationRecord
+        include EachBatch
         include FromUnion
 
         self.table_name = :batched_background_migration_jobs

lib/gitlab/database/background_migration/batched_migration.rb

@@ -68,6 +68,17 @@ module Gitlab
           )
         end
 
+        def retry_failed_jobs!
+          batched_jobs.failed.each_batch(of: 100) do |batch|
+            self.class.transaction do
+              batch.lock.each(&:split_and_retry!)
+              self.active!
+            end
+          end
+
+          self.active!
+        end
+
         def next_min_value
           last_job&.max_value&.next || min_value
         end

spec/features/admin/admin_sees_background_migrations_spec.rb

@@ -10,7 +10,7 @@ RSpec.describe "Admin > Admin sees background migrations" do
   let_it_be(:finished_migration) { create(:batched_background_migration, table_name: 'finished', status: :finished) }
 
   before_all do
-    create(:batched_background_migration_job, batched_migration: failed_migration, batch_size: 30, status: :succeeded)
+    create(:batched_background_migration_job, batched_migration: failed_migration, batch_size: 10, min_value: 6, max_value: 15, status: :failed, attempts: 3)
   end
 
   before do
@@ -53,7 +53,14 @@ RSpec.describe "Admin > Admin sees background migrations" do
     end
   end
 
-  it 'can view failed migrations' do
+  context 'when there are failed migrations' do
+    before do
+      allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
+        allow(batch_class).to receive(:next_batch).with(anything, anything, batch_min_value: 6, batch_size: 5).and_return([6, 10])
+      end
+    end
+
+    it 'can view and retry them' do
       visit admin_background_migrations_path
 
       within '#content-body' do
@@ -67,8 +74,14 @@ RSpec.describe "Admin > Admin sees background migrations" do
 
         expect(page).to have_content(failed_migration.job_class_name)
         expect(page).to have_content(failed_migration.table_name)
-      expect(page).to have_content('30.00%')
+        expect(page).to have_content('0.00%')
         expect(page).to have_content(failed_migration.status.humanize)
+
+        click_button('Retry')
+        expect(page).not_to have_content(failed_migration.job_class_name)
+        expect(page).not_to have_content(failed_migration.table_name)
+        expect(page).not_to have_content('0.00%')
+      end
     end
   end
 

spec/lib/gitlab/database/background_migration/batched_migration_spec.rb

@@ -234,6 +234,42 @@ RSpec.describe Gitlab::Database::BackgroundMigration::BatchedMigration, type: :m
     end
   end
 
+  describe '#retry_failed_jobs!' do
+    let(:batched_migration) { create(:batched_background_migration, status: 'failed') }
+
+    subject(:retry_failed_jobs) { batched_migration.retry_failed_jobs! }
+
+    context 'when there are failed migration jobs' do
+      let!(:batched_background_migration_job) { create(:batched_background_migration_job, batched_migration: batched_migration, batch_size: 10, min_value: 6, max_value: 15, status: :failed, attempts: 3) }
+
+      before do
+        allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
+          allow(batch_class).to receive(:next_batch).with(anything, anything, batch_min_value: 6, batch_size: 5).and_return([6, 10])
+        end
+      end
+
+      it 'moves the status of the migration to active' do
+        retry_failed_jobs
+
+        expect(batched_migration.status).to eql 'active'
+      end
+
+      it 'changes the number of attempts to 0' do
+        retry_failed_jobs
+
+        expect(batched_background_migration_job.reload.attempts).to be_zero
+      end
+    end
+
+    context 'when there are no failed migration jobs' do
+      it 'moves the status of the migration to active' do
+        retry_failed_jobs
+
+        expect(batched_migration.status).to eql 'active'
+      end
+    end
+  end
+
   describe '#job_class_name=' do
     it_behaves_like 'an attr_writer that demodulizes assigned class names', :job_class_name
   end

spec/requests/admin/background_migrations_controller_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Admin::BackgroundMigrationsController, :enable_admin_mode do
+  let(:admin) { create(:admin) }
+
+  before do
+    sign_in(admin)
+  end
+
+  describe 'POST #retry' do
+    let(:migration) { create(:batched_background_migration, status: 'failed') }
+
+    before do
+      create(:batched_background_migration_job, batched_migration: migration, batch_size: 10, min_value: 6, max_value: 15, status: :failed, attempts: 3)
+
+      allow_next_instance_of(Gitlab::BackgroundMigration::BatchingStrategies::PrimaryKeyBatchingStrategy) do |batch_class|
+        allow(batch_class).to receive(:next_batch).with(anything, anything, batch_min_value: 6, batch_size: 5).and_return([6, 10])
+      end
+    end
+
+    subject(:retry_migration) { post retry_admin_background_migration_path(migration) }
+
+    it 'redirects the user to the admin migrations page' do
+      retry_migration
+
+      expect(response).to redirect_to(admin_background_migrations_path)
+    end
+
+    it 'retries the migration' do
+      retry_migration
+
+      expect(migration.reload.status).to eql 'active'
+    end
+
+    context 'when the migration is not failed' do
+      let(:migration) { create(:batched_background_migration, status: 'paused') }
+
+      it 'keeps the same migration status' do
+        expect { retry_migration }.not_to change { migration.reload.status }
+      end
+    end
+  end
+end