Merge branch '338004-fix-retry-mechanism-batched' into 'master'

Fix BackgroundMigration::BackfillWorkItemTypeIdForIssues

See merge request gitlab-org/gitlab!83662

lib/gitlab/background_migration/backfill_work_item_type_id_for_issues.rb

@@ -48,7 +48,7 @@ module Gitlab
 
         begin
           update_batch(sub_batch, base_type_id)
-        rescue ActiveRecord::StatementTimeout => e
+        rescue ActiveRecord::StatementTimeout, ActiveRecord::QueryCanceled => e
           update_attempt += 1
 
           if update_attempt <= MAX_UPDATE_RETRIES

spec/lib/gitlab/background_migration/backfill_work_item_type_id_for_issues_spec.rb

@@ -48,12 +48,20 @@ RSpec.describe Gitlab::BackgroundMigration::BackfillWorkItemTypeIdForIssues do
     expect { migrate }.to change { migration.batch_metrics.timings }
   end
 
-  it 'retries on ActiveRecord::StatementTimeout' do
-    expect(migration).to receive(:update_batch).exactly(3).times.and_raise(ActiveRecord::StatementTimeout)
-    expect(migration).to receive(:sleep).with(30).twice
+  context 'when database timeouts' do
+    using RSpec::Parameterized::TableSyntax
 
-    expect do
-      migrate
-    end.to raise_error(ActiveRecord::StatementTimeout)
+    where(error_class: [ActiveRecord::StatementTimeout, ActiveRecord::QueryCanceled])
+
+    with_them do
+      it 'retries on timeout error' do
+        expect(migration).to receive(:update_batch).exactly(3).times.and_raise(error_class)
+        expect(migration).to receive(:sleep).with(30).twice
+
+        expect do
+          migrate
+        end.to raise_error(error_class)
+      end
+    end
   end
 end