diff --git a/lib/gitlab/background_migration/backfill_desired_sharding_key_job.rb b/lib/gitlab/background_migration/backfill_desired_sharding_key_job.rb new file mode 100644 index 0000000000000000000000000000000000000000..4552d03b211144eb66b1edcf3abc70046d9ddc40 --- /dev/null +++ b/lib/gitlab/background_migration/backfill_desired_sharding_key_job.rb @@ -0,0 +1,29 @@ +# 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 diff --git a/spec/lib/gitlab/background_migration/backfill_desired_sharding_key_job_spec.rb b/spec/lib/gitlab/background_migration/backfill_desired_sharding_key_job_spec.rb new file mode 100644 index 0000000000000000000000000000000000000000..24ed3ec85b8285e9d4bbbc475b66cdfdf379cb14 --- /dev/null +++ b/spec/lib/gitlab/background_migration/backfill_desired_sharding_key_job_spec.rb @@ -0,0 +1,71 @@ +# 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 diff --git a/spec/support/shared_examples/lib/gitlab/background_migration/backfill_desired_sharding_key_shared_examples.rb b/spec/support/shared_examples/lib/gitlab/background_migration/backfill_desired_sharding_key_shared_examples.rb new file mode 100644 index 0000000000000000000000000000000000000000..17d71de47301ba26a91195567a84cce55327fc0f --- /dev/null +++ b/spec/support/shared_examples/lib/gitlab/background_migration/backfill_desired_sharding_key_shared_examples.rb @@ -0,0 +1,35 @@ +# 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