Merge branch 'backfill-desired-sharding-key-job' into 'master'

Generic job for backfilling desired sharding keys

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/144401



Merged-by: Alper Akgun <aakgun@gitlab.com>
Approved-by: Dmitry Gruzd <dgruzd@gitlab.com>
Approved-by: Manoj M J <mmj@gitlab.com>
Approved-by: Alper Akgun <aakgun@gitlab.com>
Reviewed-by: Manoj M J <mmj@gitlab.com>
Co-authored-by: Tiger <twatson@gitlab.com>

lib/gitlab/background_migration/backfill_desired_sharding_key_job.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    # rubocop: disable BackgroundMigration/FeatureCategory -- Feature category to be specified by inheriting class
+    class BackfillDesiredShardingKeyJob < BatchedMigrationJob
+      job_arguments :backfill_column, :backfill_via_table, :backfill_via_column, :backfill_via_foreign_key
+
+      scope_to ->(relation) { relation.where(backfill_column => nil) }
+
+      def perform
+        each_sub_batch do |sub_batch|
+          sub_batch.connection.execute(construct_query(sub_batch: sub_batch))
+        end
+      end
+
+      def construct_query(sub_batch:)
+        <<~SQL
+          UPDATE #{batch_table}
+          SET #{backfill_column} = #{backfill_via_table}.#{backfill_via_column}
+          FROM #{backfill_via_table}
+          WHERE #{backfill_via_table}.id = #{batch_table}.#{backfill_via_foreign_key}
+          AND #{batch_table}.id IN (#{sub_batch.select(:id).to_sql})
+        SQL
+      end
+    end
+    # rubocop: enable BackgroundMigration/FeatureCategory
+  end
+end

spec/lib/gitlab/background_migration/backfill_desired_sharding_key_job_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillDesiredShardingKeyJob,
+  feature_category: :cell,
+  schema: 20231114034017 do
+  let(:example_job_class) do
+    Class.new(described_class) do
+      operation_name :backfill_merge_request_diffs_project_id
+      feature_category :cell
+    end
+  end
+
+  let!(:start_id) { table(:merge_request_diffs).minimum(:id) }
+  let!(:end_id) { table(:merge_request_diffs).maximum(:id) }
+  let!(:migration) do
+    example_job_class.new(
+      start_id: start_id,
+      end_id: end_id,
+      batch_table: :merge_request_diffs,
+      batch_column: :id,
+      sub_batch_size: 10,
+      pause_ms: 2,
+      connection: ::ApplicationRecord.connection,
+      job_arguments: [
+        :project_id,
+        :merge_requests,
+        :target_project_id,
+        :merge_request_id
+      ]
+    )
+  end
+
+  describe '#perform' do
+    let!(:diffs_without_project_id) do
+      13.times do
+        namespace = table(:namespaces).create!(name: 'my namespace', path: 'my-namespace')
+        project = table(:projects).create!(name: 'my project', path: 'my-project', namespace_id: namespace.id,
+          project_namespace_id: namespace.id)
+        merge_request = table(:merge_requests).create!(target_project_id: project.id, target_branch: 'main',
+          source_branch: 'not-main')
+        table(:merge_request_diffs).create!(merge_request_id: merge_request.id, project_id: nil)
+      end
+    end
+
+    it 'backfills the missing project_id for the batch' do
+      backfilled_diffs = table(:merge_request_diffs)
+        .joins('INNER JOIN merge_requests ON merge_request_diffs.merge_request_id = merge_requests.id')
+        .where('merge_request_diffs.project_id = merge_requests.target_project_id')
+
+      expect do
+        migration.perform
+      end.to change { backfilled_diffs.count }.from(0).to(13)
+    end
+  end
+
+  describe '#constuct_query' do
+    it 'constructs a query using the supplied job arguments' do
+      sub_batch = table(:merge_request_diffs).all
+
+      expect(migration.construct_query(sub_batch: sub_batch)).to eq(<<~SQL)
+        UPDATE merge_request_diffs
+        SET project_id = merge_requests.target_project_id
+        FROM merge_requests
+        WHERE merge_requests.id = merge_request_diffs.merge_request_id
+        AND merge_request_diffs.id IN (#{sub_batch.select(:id).to_sql})
+      SQL
+    end
+  end
+end

spec/support/shared_examples/lib/gitlab/background_migration/backfill_desired_sharding_key_shared_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'desired sharding key backfill job' do
+  let!(:connection) { table(batch_table).connection }
+  let!(:starting_id) { table(batch_table).pluck(:id).min }
+  let!(:end_id) { table(batch_table).pluck(:id).max }
+
+  let!(:migration) do
+    described_class.new(
+      start_id: starting_id,
+      end_id: end_id,
+      batch_table: batch_table,
+      batch_column: :id,
+      sub_batch_size: 10,
+      pause_ms: 2,
+      connection: connection,
+      job_arguments: [
+        backfill_column,
+        backfill_via_table,
+        backfill_via_column,
+        backfill_via_foreign_key
+      ]
+    )
+  end
+
+  it 'performs without error' do
+    expect { migration.perform }.not_to raise_error
+  end
+
+  it 'constructs a valid query' do
+    query = migration.construct_query(sub_batch: table(batch_table).all)
+
+    expect { connection.execute(query) }.not_to raise_error
+  end
+end