Use sub-batches in BackfillIntegrationsTypeNew background migration

Previously we were ignoring the `sub_batch_size` argument and using
`start_id` and `stop_id` directly, which led to timeouts on staging.

The migration succeeded on production, but this change should help
increase the chances for the migration to succeed on self-managed
instances as well.

Changelog: fixed

lib/gitlab/background_migration/backfill_integrations_type_new.rb

@@ -6,8 +6,32 @@ module Gitlab
     # the real class name, rather than the legacy class name in `type`
     # which is mapped via `Gitlab::Integrations::StiType`.
     class BackfillIntegrationsTypeNew
-      def perform(start_id, stop_id, *args)
-        ActiveRecord::Base.connection.execute(<<~SQL)
+      include Gitlab::Database::DynamicModelHelpers
+
+      def perform(start_id, stop_id, batch_table, batch_column, sub_batch_size, pause_ms)
+        parent_batch_relation = define_batchable_model(batch_table)
+          .where(batch_column => start_id..stop_id)
+
+        parent_batch_relation.each_batch(column: batch_column, of: sub_batch_size) do |sub_batch|
+          process_sub_batch(sub_batch)
+
+          sleep(pause_ms * 0.001) if pause_ms > 0
+        end
+      end
+
+      private
+
+      def connection
+        ActiveRecord::Base.connection
+      end
+
+      def process_sub_batch(sub_batch)
+        # Extract the start/stop IDs from the current sub-batch
+        sub_start_id, sub_stop_id = sub_batch.pluck(Arel.sql('MIN(id), MAX(id)')).first
+
+        # This matches the mapping from the INSERT trigger added in
+        # db/migrate/20210721135638_add_triggers_to_integrations_type_new.rb
+        connection.execute(<<~SQL)
           WITH mapping(old_type, new_type) AS (VALUES
             ('AsanaService',                   'Integrations::Asana'),
             ('AssemblaService',                'Integrations::Assembla'),
@@ -53,7 +77,7 @@ module Gitlab
 
           UPDATE integrations SET type_new = mapping.new_type
           FROM mapping
-          WHERE integrations.id BETWEEN #{start_id} AND #{stop_id}
+          WHERE integrations.id BETWEEN #{sub_start_id} AND #{sub_stop_id}
             AND integrations.type = mapping.old_type
         SQL
       end

spec/lib/gitlab/background_migration/backfill_integrations_type_new_spec.rb

@@ -3,6 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::BackgroundMigration::BackfillIntegrationsTypeNew do
+  let(:migration) { described_class.new }
   let(:integrations) { table(:integrations) }
   let(:namespaced_integrations) { Gitlab::Integrations::StiType.namespaced_integrations }
 
@@ -12,12 +13,31 @@ RSpec.describe Gitlab::BackgroundMigration::BackfillIntegrationsTypeNew do
     namespaced_integrations.each_with_index do |type, i|
       integrations.create!(id: i + 1, type: "#{type}Service")
     end
+
+    integrations.create!(id: namespaced_integrations.size + 1, type: 'LegacyService')
   ensure
     integrations.connection.execute 'ALTER TABLE integrations ENABLE TRIGGER "trigger_type_new_on_insert"'
   end
 
   it 'backfills `type_new` for the selected records' do
-    described_class.new.perform(2, 10)
+    # We don't want to mock `Kernel.sleep`, so instead we mock it on the migration
+    # class before it gets forwarded.
+    expect(migration).to receive(:sleep).with(0.05).exactly(5).times
+
+    queries = ActiveRecord::QueryRecorder.new do
+      migration.perform(2, 10, :integrations, :id, 2, 50)
+    end
+
+    expect(queries.count).to be(16)
+    expect(queries.log.grep(/^SELECT/).size).to be(11)
+    expect(queries.log.grep(/^UPDATE/).size).to be(5)
+    expect(queries.log.grep(/^UPDATE/).join.scan(/WHERE .*/)).to eq([
+      'WHERE integrations.id BETWEEN 2 AND 3',
+      'WHERE integrations.id BETWEEN 4 AND 5',
+      'WHERE integrations.id BETWEEN 6 AND 7',
+      'WHERE integrations.id BETWEEN 8 AND 9',
+      'WHERE integrations.id BETWEEN 10 AND 10'
+    ])
 
     expect(integrations.where(id: 2..10).pluck(:type, :type_new)).to contain_exactly(
       ['AssemblaService',           'Integrations::Assembla'],