git.cworth.org Git - sup/blob - bin/sup-import

   1 #!/usr/bin/env ruby
   2
   3 require 'rubygems'
   4 require 'highline'
   5 require "sup"
   6
   7 class Float
   8   def to_s; sprintf '%.2f', self; end
   9 end
  10
  11 class Numeric
  12   def to_time_s
  13     i = to_i
  14     sprintf "%d:%02d:%02d", i / 3600, (i / 60) % 60, i % 60
  15   end
  16 end
  17
  18 def time
  19   startt = Time.now
  20   yield
  21   Time.now - startt
  22 end
  23
  24 def educate_user
  25   $stderr.puts <<EOS
  26 Loads messages into the Sup index, adding sources as needed to the
  27 source list.
  28
  29 Usage:
  30   sup-import [options] <source>*
  31 where <source>* is zero or more source descriptions (e.g., mbox
  32 filenames on disk, or imap/imaps URIs).
  33
  34 If the sources listed are not already in the Sup source list,
  35 they will be added to it, as parameterized by the following options:
  36   --archive: messages from these sources will not appear in the inbox
  37   --unusual: these sources will not be polled when the flag --the-usual
  38              is called
  39
  40 Regardless of whether the sources are new or not, they will be polled,
  41 and any new messages will be added to the index, as parameterized by
  42 the following options:
  43   --force-archive: regardless of the source "archive" flag, any new
  44                    messages found will not appear in the inbox.
  45   --force-read:    any messages found will not be marked as new.
  46
  47 The following options can also be specified:
  48   --the-usual:     import new messages from all usual sources
  49   --rebuild:       rebuild the index for the specified sources rather than
  50                    just adding new messages. Useful if the sources
  51                    have changed in any way *other* than new messages
  52                    being added.
  53   --force-rebuild: force a rebuild of all messages in the inbox, not just
  54                    ones that have changed. You probably won't need this
  55                    unless William changes the index format.
  56   --optimize:      optimize the index after adding any new messages.
  57   --help:          don't do anything, just show this message.
  58 EOS
  59   exit
  60 end
  61
  62 educate_user if ARGV.member? '--help'
  63
  64 archive = ARGV.delete "--archive"
  65 unusual = ARGV.delete "--unusual"
  66 force_archive = ARGV.delete "--force-archive"
  67 force_read = ARGV.delete "--force-read"
  68 the_usual = ARGV.delete "--the-usual"
  69 rebuild = ARGV.delete "--rebuild"
  70 force_rebuild = ARGV.delete "--force-rebuild"
  71 optimize = ARGV.delete "--optimize"
  72 start_at = # ok really need to use optparse or something now
  73   if(i = ARGV.index("--start-at"))
  74     raise "start-at requires a numeric argument: #{ARGV[i + 1].inspect}" unless ARGV.length > (i + 1) && ARGV[i + 1] =~ /\d/
  75     ARGV.delete_at i
  76     ARGV.delete_at(i).to_i # whoa!
  77   end
  78
  79 if(o = ARGV.find { |x| x =~ /^--/ })
  80   $stderr.puts "error: unknown option #{o}"
  81   educate_user
  82 end
  83
  84 Redwood::start
  85
  86 index = Redwood::Index.new
  87 index.load
  88
  89 h = HighLine.new
  90
  91 sources = ARGV.map do |fn|
  92   fn = "mbox://#{fn}" unless fn =~ %r!://!
  93   source = index.source_for fn
  94   unless source
  95     source =
  96       case fn
  97       when %r!^mbox\+ssh://!
  98         username = h.ask("Username for #{fn}: ");
  99         password = h.ask("Password for #{fn}: ") { |q| q.echo = false }
 100         puts # why?
 101         Redwood::MBox::SSHLoader.new(fn, username, password, nil, !unusual, !!archive)
 102       when %r!^imaps?://!
 103         username = h.ask("Username for #{fn}: ");
 104         password = h.ask("Password for #{fn}: ") { |q| q.echo = false }
 105         puts # why?
 106         Redwood::IMAP.new(fn, username, password, nil, !unusual, !!archive)
 107       else
 108         Redwood::MBox::Loader.new(fn, nil, !unusual, !!archive)
 109       end
 110     index.add_source source
 111   end
 112   source
 113 end
 114
 115 sources = (sources + index.usual_sources).uniq if the_usual
 116 if rebuild || force_rebuild
 117   if start_at
 118     sources.each { |s| s.seek_to! start_at }
 119   else
 120     sources.each { |s| s.reset! }
 121   end
 122 end
 123
 124 found = {}
 125 start = Time.now
 126 begin
 127   sources.each do |source|
 128     if source.broken?
 129       puts "error loading messages from #{source}: #{source.broken_msg}"
 130       next
 131     end
 132     next if source.done?
 133     puts "loading from #{source}... "
 134     num = 0
 135     start_offset = nil
 136     source.each do |offset, labels|
 137       start_offset ||= offset
 138       labels -= [:inbox] if force_archive
 139       labels -= [:unread] if force_read
 140       begin
 141         m = Redwood::Message.new :source => source, :source_info => offset, :labels => labels
 142         if found[m.id]
 143           puts "skipping duplicate message #{m.id}"
 144           next
 145         else
 146           found[m.id] = true
 147         end
 148         m.remove_label :unread if m.status == "RO" unless force_read
 149         puts "# message at #{offset}, labels: #{labels * ', '}" unless rebuild || force_rebuild
 150         if (rebuild || force_rebuild) &&
 151             (docid, entry = index.load_entry_for_id(m.id)) && entry
 152           if force_rebuild || entry[:source_info].to_i != offset
 153             puts "replacing message #{m.id} labels #{entry[:label].inspect} (offset #{entry[:source_info]} => #{offset})"
 154             m.labels = entry[:label].split.map { |l| l.intern }
 155             num += 1 if index.update_message m, source, offset
 156           end
 157         else
 158           num += 1 if index.add_message m
 159         end
 160       rescue Redwood::MessageFormatError, Redwood::SourceError => e
 161         $stderr.puts "ignoring erroneous message at #{source}##{offset}: #{e.message}"
 162       end
 163       if num % 1000 == 0 && num > 0
 164         elapsed = Time.now - start
 165         pctdone = (offset.to_f - start_offset) / (source.total.to_f - start_offset)
 166         remaining = (source.total.to_f - offset.to_f) * (elapsed.to_f / (offset.to_f - start_offset))
 167         puts "## #{num} (#{(pctdone * 100.0)}% done) read; #{elapsed.to_time_s} elapsed; est. #{remaining.to_time_s} remaining"
 168       end
 169     end
 170     puts "loaded #{num} messages" unless num == 0
 171   end
 172 ensure
 173   index.save
 174   Redwood::finish
 175 end
 176
 177 if rebuild || force_rebuild
 178   puts "deleting missing messages from the index..."
 179   numdel = num = 0
 180   sources.each do |source|
 181     raise "no source id for #{source}" unless source.id
 182     q = "+source_id:#{source.id}"
 183     q += " +source_info: >= #{start_at}" if start_at
 184     #p q
 185     num += index.index.search_each(q, :limit => :all) do |docid, score|
 186       mid = index.index[docid][:message_id]
 187       next if found[mid]
 188       puts "deleting #{mid}"
 189       index.index.delete docid
 190       numdel += 1
 191     end
 192     #p num
 193   end
 194   puts "deleted #{numdel} / #{num} messages"
 195 end
 196
 197 if optimize
 198   puts "optimizing index..."
 199   optt = time { index.index.optimize }
 200   puts "optimized index of size #{index.size} in #{optt}s."
 201 end