gitlab-foss/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

      def perform
        each_sub_batch do |sub_batch|
          sub_batch.connection.execute(construct_query(sub_batch: sub_batch.where(backfill_column => nil)))
        end
      end

      def backfill_via_table_primary_key
        'id'
      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}.#{backfill_via_table_primary_key} = #{batch_table}.#{backfill_via_foreign_key}
          AND #{batch_table}.#{batch_column} IN (#{sub_batch.select(batch_column).to_sql})
        SQL
      end
    end
    # rubocop: enable BackgroundMigration/FeatureCategory
  end
end