git.notmuchmail.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 puts "loading index..."
  87 index = Redwood::Index.new
  88 index.load
  89 puts "loaded index of #{index.size} messages"
  90
  91 h = HighLine.new
  92
  93 sources = ARGV.map do |fn|
  94   fn = "mbox://#{fn}" unless fn =~ %r!://!
  95   source = index.source_for fn
  96   unless source
  97     source =
  98       case fn
  99       when %r!^mbox\+ssh://!
 100         username = h.ask("Username for #{fn}: ");
 101         password = h.ask("Password for #{fn}: ") { |q| q.echo = false }
 102         puts # why?
 103         Redwood::MBox::SSHLoader.new(fn, username, password, nil, !unusual, !!archive)
 104       when %r!^imaps?://!
 105         username = h.ask("Username for #{fn}: ");
 106         password = h.ask("Password for #{fn}: ") { |q| q.echo = false }
 107         puts # why?
 108         Redwood::IMAP.new(fn, username, password, nil, !unusual, !!archive)
 109       else
 110         Redwood::MBox::Loader.new(fn, nil, !unusual, !!archive)
 111       end
 112     index.add_source source
 113   end
 114   source
 115 end
 116
 117 sources = (sources + index.usual_sources).uniq if the_usual
 118 if rebuild || force_rebuild
 119   if start_at
 120     sources.each { |s| s.seek_to! start_at }
 121   else
 122     sources.each { |s| s.reset! }
 123   end
 124 end
 125
 126 found = {}
 127 start = Time.now
 128 begin
 129   sources.each do |source|
 130     if source.broken?
 131       puts "error loading messages from #{source}: #{source.broken_msg}"
 132       next
 133     end
 134     next if source.done?
 135     puts "loading from #{source}... "
 136     num = 0
 137     start_offset = nil
 138     source.each do |offset, labels|
 139       start_offset ||= offset
 140       labels -= [:inbox] if force_archive
 141       labels -= [:unread] if force_read
 142       begin
 143         m = Redwood::Message.new :source => source, :source_info => offset, :labels => labels
 144         if found[m.id]
 145           puts "skipping duplicate message #{m.id}"
 146           next
 147         else
 148           found[m.id] = true
 149         end
 150         m.remove_label :unread if m.status == "RO" unless force_read
 151         puts "# message at #{offset}, labels: #{labels * ', '}" unless rebuild || force_rebuild
 152         if (rebuild || force_rebuild) &&
 153             (docid, entry = index.load_entry_for_id(m.id)) && entry
 154           if force_rebuild || entry[:source_info].to_i != offset
 155             puts "replacing message #{m.id} labels #{entry[:label].inspect} (offset #{entry[:source_info]} => #{offset})"
 156             m.labels = entry[:label].split.map { |l| l.intern }
 157             num += 1 if index.update_message m, source, offset
 158           end
 159         else
 160           num += 1 if index.add_message m
 161         end
 162       rescue Redwood::MessageFormatError, Redwood::SourceError => e
 163         $stderr.puts "ignoring erroneous message at #{source}##{offset}: #{e.message}"
 164       end
 165       if num % 1000 == 0 && num > 0
 166         elapsed = Time.now - start
 167         pctdone = (offset.to_f - start_offset) / (source.total.to_f - start_offset)
 168         remaining = (source.total.to_f - offset.to_f) * (elapsed.to_f / (offset.to_f - start_offset))
 169         puts "## #{num} (#{(pctdone * 100.0)}% done) read; #{elapsed.to_time_s} elapsed; est. #{remaining.to_time_s} remaining"
 170       end
 171     end
 172     puts "loaded #{num} messages" unless num == 0
 173   end
 174 ensure
 175   index.save
 176   Redwood::finish
 177 end
 178
 179 if rebuild || force_rebuild
 180   puts "deleting missing messages from the index..."
 181   numdel = num = 0
 182   sources.each do |source|
 183     raise "no source id for #{source}" unless source.id
 184     q = "+source_id:#{source.id}"
 185     q += " +source_info: >= #{start_at}" if start_at
 186     #p q
 187     num += index.index.search_each(q, :limit => :all) do |docid, score|
 188       mid = index.index[docid][:message_id]
 189       next if found[mid]
 190       puts "deleting #{mid}"
 191       index.index.delete docid
 192       numdel += 1
 193     end
 194     #p num
 195   end
 196   puts "deleted #{numdel} / #{num} messages"
 197 end
 198
 199 if optimize
 200   puts "optimizing index..."
 201   optt = time { index.index.optimize }
 202   puts "optimized index of size #{index.size} in #{optt}s."
 203 end