2018-08-26 23:53:06 +09:00
|
|
|
# frozen_string_literal: true
|
|
|
|
|
|
|
|
dev_null = Logger.new('/dev/null')
|
|
|
|
|
|
|
|
Rails.logger = dev_null
|
|
|
|
ActiveRecord::Base.logger = dev_null
|
2018-09-09 20:33:36 +09:00
|
|
|
ActiveJob::Base.logger = dev_null
|
2018-08-26 23:53:06 +09:00
|
|
|
HttpLog.configuration.logger = dev_null
|
|
|
|
Paperclip.options[:log] = false
|
2020-07-15 01:10:35 +09:00
|
|
|
Chewy.logger = dev_null
|
2019-09-10 20:48:48 +09:00
|
|
|
|
2023-05-23 23:08:26 +09:00
|
|
|
module Mastodon::CLI
|
|
|
|
module Helper
|
2020-04-27 06:29:08 +09:00
|
|
|
def dry_run?
|
|
|
|
options[:dry_run]
|
|
|
|
end
|
|
|
|
|
2019-09-10 20:48:48 +09:00
|
|
|
def create_progress_bar(total = nil)
|
|
|
|
ProgressBar.create(total: total, format: '%c/%u |%b%i| %e')
|
|
|
|
end
|
|
|
|
|
2022-04-30 05:43:07 +09:00
|
|
|
def reset_connection_pools!
|
|
|
|
ActiveRecord::Base.establish_connection(ActiveRecord::Base.configurations[Rails.env].dup.tap { |config| config['pool'] = options[:concurrency] + 1 })
|
|
|
|
RedisConfiguration.establish_pool(options[:concurrency])
|
|
|
|
end
|
|
|
|
|
2019-09-10 20:48:48 +09:00
|
|
|
def parallelize_with_progress(scope)
|
2019-10-07 11:24:05 +09:00
|
|
|
if options[:concurrency] < 1
|
|
|
|
say('Cannot run with this concurrency setting, must be at least 1', :red)
|
|
|
|
exit(1)
|
|
|
|
end
|
|
|
|
|
2022-04-30 05:43:07 +09:00
|
|
|
reset_connection_pools!
|
2019-09-10 20:48:48 +09:00
|
|
|
|
|
|
|
progress = create_progress_bar(scope.count)
|
|
|
|
pool = Concurrent::FixedThreadPool.new(options[:concurrency])
|
|
|
|
total = Concurrent::AtomicFixnum.new(0)
|
|
|
|
aggregate = Concurrent::AtomicFixnum.new(0)
|
|
|
|
|
|
|
|
scope.reorder(nil).find_in_batches do |items|
|
|
|
|
futures = []
|
|
|
|
|
|
|
|
items.each do |item|
|
|
|
|
futures << Concurrent::Future.execute(executor: pool) do
|
2023-02-19 07:09:40 +09:00
|
|
|
if !progress.total.nil? && progress.progress + 1 > progress.total
|
|
|
|
# The number of items has changed between start and now,
|
|
|
|
# since there is no good way to predict the final count from
|
|
|
|
# here, just change the progress bar to an indeterminate one
|
|
|
|
|
|
|
|
progress.total = nil
|
|
|
|
end
|
|
|
|
|
|
|
|
progress.log("Processing #{item.id}") if options[:verbose]
|
|
|
|
|
2023-03-21 04:02:58 +09:00
|
|
|
Chewy.strategy(:mastodon) do
|
|
|
|
result = ActiveRecord::Base.connection_pool.with_connection do
|
|
|
|
yield(item)
|
|
|
|
ensure
|
|
|
|
RedisConfiguration.pool.checkin if Thread.current[:redis]
|
|
|
|
Thread.current[:redis] = nil
|
|
|
|
end
|
|
|
|
|
|
|
|
aggregate.increment(result) if result.is_a?(Integer)
|
2019-09-10 20:48:48 +09:00
|
|
|
end
|
2023-02-19 07:09:40 +09:00
|
|
|
rescue => e
|
|
|
|
progress.log pastel.red("Error processing #{item.id}: #{e}")
|
|
|
|
ensure
|
|
|
|
progress.increment
|
2019-09-10 20:48:48 +09:00
|
|
|
end
|
|
|
|
end
|
|
|
|
|
|
|
|
total.increment(items.size)
|
|
|
|
futures.map(&:value)
|
|
|
|
end
|
|
|
|
|
2019-10-07 11:24:05 +09:00
|
|
|
progress.stop
|
2019-09-10 20:48:48 +09:00
|
|
|
|
|
|
|
[total.value, aggregate.value]
|
|
|
|
end
|
|
|
|
|
|
|
|
def pastel
|
|
|
|
@pastel ||= Pastel.new
|
|
|
|
end
|
|
|
|
end
|
|
|
|
end
|