git.cworth.org Git - sup/blob - bin/sup-tweak-labels

   1 #!/usr/bin/env ruby
   2
   3 require 'rubygems'
   4 require 'trollop'
   5 require 'enumerator'
   6 require "sup"
   7
   8 class Float
   9   def to_s; sprintf '%.2f', self; end
  10    def to_time_s
  11      infinite? ? "unknown" : super
  12    end
  13 end
  14
  15 class Numeric
  16   def to_time_s
  17     i = to_i
  18     sprintf "%d:%02d:%02d", i / 3600, (i / 60) % 60, i % 60
  19   end
  20 end
  21
  22 def time
  23   startt = Time.now
  24   yield
  25   Time.now - startt
  26 end
  27
  28 opts = Trollop::options do
  29   version "sup-tweak-labels (sup #{Redwood::VERSION})"
  30   banner <<EOS
  31 Batch modification of message state for messages already in the index.
  32
  33 Usage:
  34   sup-tweak-labels [options] <source>*
  35
  36 where <source>* is zero or more source URIs. Supported source URI schemes can
  37 be seen by running "sup-add --help".
  38
  39 Options:
  40 EOS
  41   opt :add, "One or more labels (comma-separated) to add to every message from the specified sources", :default => ""
  42   opt :remove, "One or more labels (comma-separated) to remove from every message from the specified sources, if those labels are present", :default => ""
  43   opt :query, "A Sup search query", :type => String
  44
  45   text <<EOS
  46
  47 Other options:
  48 EOS
  49   opt :verbose, "Print message ids as they're processed."
  50   opt :very_verbose, "Print message names and subjects as they're processed."
  51   opt :all_sources, "Scan over all sources.", :short => :none
  52   opt :dry_run, "Don't actually modify the index. Probably only useful with --verbose.", :short => "-n"
  53   opt :version, "Show version information", :short => :none
  54 end
  55 opts[:verbose] = true if opts[:very_verbose]
  56
  57 add_labels = opts[:add].to_set_of_symbols ","
  58 remove_labels = opts[:remove].to_set_of_symbols ","
  59
  60 Trollop::die "nothing to do: no labels to add or remove" if add_labels.empty? && remove_labels.empty?
  61 Trollop::die "no sources specified" if ARGV.empty?
  62
  63 Redwood::start
  64 index = Redwood::Index.init
  65 index.lock_interactively or exit
  66 begin
  67   index.load
  68
  69   source_ids = if opts[:all_sources]
  70     Redwood::SourceManager.sources
  71   else
  72     ARGV.map do |uri|
  73       Redwood::SourceManager.source_for uri or Trollop::die "Unknown source: #{uri}. Did you add it with sup-add first?"
  74     end
  75   end.map { |s| s.id }
  76   Trollop::die "nothing to do: no sources" if source_ids.empty?
  77
  78   query = "+(" + source_ids.map { |id| "source_id:#{id}" }.join(" OR ") + ")"
  79   if add_labels.empty?
  80     ## if all we're doing is removing labels, we can further restrict the
  81     ## query to only messages with those labels
  82     query += " +(" + remove_labels.map { |l| "label:#{l}" }.join(" ") + ")"
  83   end
  84   query += ' ' + opts[:query] if opts[:query]
  85
  86   parsed_query = index.parse_query query
  87   parsed_query.merge! :load_spam => true, :load_deleted => true, :load_killed => true
  88   ids = Enumerable::Enumerator.new(index, :each_id, parsed_query).map
  89   num_total = ids.size
  90
  91   $stderr.puts "Found #{num_total} documents across #{source_ids.length} sources. Scanning..."
  92
  93   num_changed = num_scanned = 0
  94   last_info_time = start_time = Time.now
  95   ids.each do |id|
  96     num_scanned += 1
  97
  98     m = index.build_message id
  99     old_labels = m.labels.dup
 100
 101     m.labels += add_labels
 102     m.labels -= remove_labels
 103
 104     unless m.labels == old_labels
 105       num_changed += 1
 106       puts "From #{m.from}, subject: #{m.subj}" if opts[:very_verbose]
 107       puts "#{m.id}: {#{old_labels.to_a.join ','}} => {#{m.labels.to_a.join ','}}" if opts[:verbose]
 108       puts if opts[:very_verbose]
 109       index.update_message_state m unless opts[:dry_run]
 110     end
 111
 112     if Time.now - last_info_time > 60
 113       last_info_time = Time.now
 114       elapsed = last_info_time - start_time
 115       pctdone = 100.0 * num_scanned.to_f / num_total.to_f
 116       remaining = (100.0 - pctdone) * (elapsed.to_f / pctdone)
 117       $stderr.puts "## #{num_scanned} (#{pctdone}%) read; #{elapsed.to_time_s} elapsed; #{remaining.to_time_s} remaining"
 118     end
 119   end
 120   $stderr.puts "Scanned #{num_scanned} / #{num_total} messages and changed #{num_changed}."
 121
 122   unless num_changed == 0
 123     $stderr.puts "Optimizing index..."
 124     index.optimize unless opts[:dry_run]
 125   end
 126
 127 rescue Exception => e
 128   File.open("sup-exception-log.txt", "w") { |f| f.puts e.backtrace }
 129   raise
 130 ensure
 131   index.save
 132   Redwood::finish
 133   index.unlock
 134 end
 135