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

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