lib/mastodon/cli_helper.rb

   1 # frozen_string_literal: true
   2
   3 dev_null = Logger.new('/dev/null')
   4
   5 Rails.logger                 = dev_null
   6 ActiveRecord::Base.logger    = dev_null
   7 ActiveJob::Base.logger       = dev_null
   8 HttpLog.configuration.logger = dev_null
   9 Paperclip.options[:log]      = false
  10 Chewy.logger                 = dev_null
  11
  12 module Mastodon
  13   module CLIHelper
  14     def dry_run?
  15       options[:dry_run]
  16     end
  17
  18     def create_progress_bar(total = nil)
  19       ProgressBar.create(total: total, format: '%c/%u |%b%i| %e')
  20     end
  21
  22     def parallelize_with_progress(scope)
  23       if options[:concurrency] < 1
  24         say('Cannot run with this concurrency setting, must be at least 1', :red)
  25         exit(1)
  26       end
  27
  28       ActiveRecord::Base.configurations[Rails.env]['pool'] = options[:concurrency] + 1
  29
  30       progress  = create_progress_bar(scope.count)
  31       pool      = Concurrent::FixedThreadPool.new(options[:concurrency])
  32       total     = Concurrent::AtomicFixnum.new(0)
  33       aggregate = Concurrent::AtomicFixnum.new(0)
  34
  35       scope.reorder(nil).find_in_batches do |items|
  36         futures = []
  37
  38         items.each do |item|
  39           futures << Concurrent::Future.execute(executor: pool) do
  40             begin
  41               if !progress.total.nil? && progress.progress + 1 > progress.total
  42                 # The number of items has changed between start and now,
  43                 # since there is no good way to predict the final count from
  44                 # here, just change the progress bar to an indeterminate one
  45
  46                 progress.total = nil
  47               end
  48
  49               progress.log("Processing #{item.id}") if options[:verbose]
  50
  51               result = ActiveRecord::Base.connection_pool.with_connection do
  52                 yield(item)
  53               end
  54
  55               aggregate.increment(result) if result.is_a?(Integer)
  56             rescue => e
  57               progress.log pastel.red("Error processing #{item.id}: #{e}")
  58             ensure
  59               progress.increment
  60             end
  61           end
  62         end
  63
  64         total.increment(items.size)
  65         futures.map(&:value)
  66       end
  67
  68       progress.stop
  69
  70       [total.value, aggregate.value]
  71     end
  72
  73     def pastel
  74       @pastel ||= Pastel.new
  75     end
  76   end
  77 end