69 lines
2.3 KiB
Ruby
69 lines
2.3 KiB
Ruby
# frozen_string_literal: true
|
|
|
|
module Gitlab
|
|
module BackgroundMigration
|
|
# Base class for batched background migrations. Subclasses should implement the `#perform`
|
|
# method as the entry point for the job's execution, which will be called with the migration
|
|
# arguments (if any).
|
|
class BatchedMigrationJob
|
|
include Gitlab::Database::DynamicModelHelpers
|
|
|
|
def initialize(start_id:, end_id:, batch_table:, batch_column:, sub_batch_size:, pause_ms:, connection:)
|
|
@start_id = start_id
|
|
@end_id = end_id
|
|
@batch_table = batch_table
|
|
@batch_column = batch_column
|
|
@sub_batch_size = sub_batch_size
|
|
@pause_ms = pause_ms
|
|
@connection = connection
|
|
end
|
|
|
|
def perform(*job_arguments)
|
|
raise NotImplementedError, "subclasses of #{self.class.name} must implement #{__method__}"
|
|
end
|
|
|
|
def batch_metrics
|
|
@batch_metrics ||= Gitlab::Database::BackgroundMigration::BatchMetrics.new
|
|
end
|
|
|
|
private
|
|
|
|
attr_reader :start_id, :end_id, :batch_table, :batch_column, :sub_batch_size, :pause_ms, :connection
|
|
|
|
def each_sub_batch(operation_name: :default, batching_arguments: {}, batching_scope: nil)
|
|
all_batching_arguments = { column: batch_column, of: sub_batch_size }.merge(batching_arguments)
|
|
|
|
parent_relation = parent_batch_relation(batching_scope)
|
|
|
|
parent_relation.each_batch(**all_batching_arguments) do |relation|
|
|
batch_metrics.instrument_operation(operation_name) do
|
|
yield relation
|
|
end
|
|
|
|
sleep([pause_ms, 0].max * 0.001)
|
|
end
|
|
end
|
|
|
|
def distinct_each_batch(operation_name: :default, batching_arguments: {})
|
|
all_batching_arguments = { column: batch_column, of: sub_batch_size }.merge(batching_arguments)
|
|
|
|
parent_batch_relation.distinct_each_batch(**all_batching_arguments) do |relation|
|
|
batch_metrics.instrument_operation(operation_name) do
|
|
yield relation
|
|
end
|
|
|
|
sleep([pause_ms, 0].max * 0.001)
|
|
end
|
|
end
|
|
|
|
def parent_batch_relation(batching_scope = nil)
|
|
parent_relation = define_batchable_model(batch_table, connection: connection)
|
|
.where(batch_column => start_id..end_id)
|
|
|
|
return parent_relation unless batching_scope
|
|
|
|
batching_scope.call(parent_relation)
|
|
end
|
|
end
|
|
end
|
|
end
|