Fixes an error on BackgroundMigration::BackfillWorkItemTypeIdForIssues

A retry mechanism is in place as we expect queries to timeout.
The wrong exception was rescued. Before we rescued only from
ActiveRecord::StatementTimeout and now we also rescue
ActiveRecord::QueryCanceled

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)
+  context 'when database timeouts' do
+    using RSpec::Parameterized::TableSyntax
+
+    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(ActiveRecord::StatementTimeout)
+        end.to raise_error(error_class)
+      end
+    end
   end
 end