mastodon/lib/active_record/batches.rb

45 lines
997 B
Ruby
Raw Normal View History

# frozen_string_literal: true
module ActiveRecord
module Batches
def pluck_each(*column_names)
relation = self
options = column_names.extract_options!
flatten = column_names.size == 1
batch_limit = options[:batch_limit] || 1_000
order = options[:order] || :asc
column_names.unshift(primary_key)
relation = relation.reorder(batch_order(order)).limit(batch_limit)
relation.skip_query_cache!
batch_relation = relation
loop do
batch = batch_relation.pluck(*column_names)
break if batch.empty?
primary_key_offset = batch.last[0]
batch.each do |record|
if flatten
yield record[1]
else
yield record[1..]
end
end
break if batch.size < batch_limit
batch_relation = relation.where(
predicate_builder[primary_key, primary_key_offset, order == :desc ? :lt : :gt]
)
end
end
end
end