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 Thread.abort_on_exception = true # make debugging possible
   8
   9 class Float
  10   def to_s; sprintf '%.2f', self; end
  11 end
  12
  13 class Numeric
  14   def to_time_s
  15     i = to_i
  16     sprintf "%d:%02d:%02d", i / 3600, (i / 60) % 60, i % 60
  17   end
  18 end
  19
  20 def time
  21   startt = Time.now
  22   yield
  23   Time.now - startt
  24 end
  25
  26 def educate_user
  27   $stderr.puts <<EOS
  28 Loads messages into the Sup index, adding sources as needed to the
  29 source list.
  30
  31 Usage:
  32   sup-import [options] <source>*
  33 where <source>* is zero or more source descriptions (e.g., mbox
  34 filenames on disk, or imap/imaps URIs).
  35
  36 If the sources listed are not already in the Sup source list,
  37 they will be added to it, as parameterized by the following options:
  38   --archive: messages from these sources will not appear in the inbox
  39   --unusual: these sources will not be polled when the flag --the-usual
  40              is called
  41
  42 Regardless of whether the sources are new or not, they will be polled,
  43 and any new messages will be added to the index, as parameterized by
  44 the following options:
  45   --force-archive: regardless of the source "archive" flag, any new
  46                    messages found will not appear in the inbox.
  47   --force-read:    any messages found will not be marked as new.
  48
  49 The following options can also be specified:
  50   --the-usual:     import new messages from all usual sources
  51   --rebuild:       rebuild the index for the specified sources rather than
  52                    just adding new messages. Useful if the sources
  53                    have changed in any way *other* than new messages
  54                    being added.
  55   --force-rebuild: force a rebuild of all messages in the inbox, not just
  56                    ones that have changed. You probably won't need this
  57                    unless William changes the index format.
  58   --optimize:      optimize the index after adding any new messages.
  59   --help:          don't do anything, just show this message.
  60 EOS
  61   exit
  62 end
  63
  64 educate_user if ARGV.member? '--help'
  65
  66 archive = ARGV.delete "--archive"
  67 unusual = ARGV.delete "--unusual"
  68 force_archive = ARGV.delete "--force-archive"
  69 force_read = ARGV.delete "--force-read"
  70 the_usual = ARGV.delete "--the-usual"
  71 rebuild = ARGV.delete "--rebuild"
  72 force_rebuild = ARGV.delete "--force-rebuild"
  73 optimize = ARGV.delete "--optimize"
  74 start_at = # ok really need to use optparse or something now
  75   if(i = ARGV.index("--start-at"))
  76     raise "start-at requires a numeric argument: #{ARGV[i + 1].inspect}" unless ARGV.length > (i + 1) && ARGV[i + 1] =~ /\d/
  77     ARGV.delete_at i
  78     ARGV.delete_at(i).to_i # whoa!
  79   end
  80
  81 if(o = ARGV.find { |x| x =~ /^--/ })
  82   $stderr.puts "error: unknown option #{o}"
  83   educate_user
  84 end
  85
  86 Redwood::start
  87
  88 index = Redwood::Index.new
  89 index.load
  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
 151         m.remove_label :unread if m.status == "RO" unless force_read
 152         puts "# message at #{offset}, labels: #{labels * ', '}"
 153         if (rebuild || force_rebuild) &&
 154             (docid, entry = index.load_entry_for_id(m.id)) && entry
 155           if force_rebuild || entry[:source_info].to_i != offset
 156             puts "replacing message #{m.id} labels #{entry[:label].inspect} (offset #{entry[:source_info]} => #{offset})"
 157             m.labels = entry[:label].split.map { |l| l.intern }
 158             num += 1 if index.update_message m, source, offset
 159           end
 160         else
 161           num += 1 if index.add_message m
 162         end
 163       rescue Redwood::MessageFormatError, Redwood::SourceError => e
 164         $stderr.puts "ignoring erroneous message at #{source}##{offset}: #{e.message}"
 165       end
 166       if num % 1000 == 0 && num > 0
 167         elapsed = Time.now - start
 168         pctdone = (offset.to_f - start_offset) / (source.total.to_f - start_offset)
 169         remaining = (source.total.to_f - offset.to_f) * (elapsed.to_f / (offset.to_f - start_offset))
 170         puts "## #{num} (#{(pctdone * 100.0)}% done) read; #{elapsed.to_time_s} elapsed; est. #{remaining.to_time_s} remaining"
 171       end
 172     end
 173     puts "loaded #{num} messages" unless num == 0
 174   end
 175 ensure
 176   $stderr.puts "saving index and sources..."
 177   index.save
 178   Redwood::finish
 179 end
 180
 181 if rebuild || force_rebuild
 182   puts "deleting missing messages from the index..."
 183   numdel = num = 0
 184   sources.each do |source|
 185     raise "no source id for #{source}" unless source.id
 186     q = "+source_id:#{source.id}"
 187     q += " +source_info: >= #{start_at}" if start_at
 188     #p q
 189     num += index.index.search_each(q, :limit => :all) do |docid, score|
 190       mid = index.index[docid][:message_id]
 191       next if found[mid]
 192       puts "deleting #{mid}"
 193       index.index.delete docid
 194       numdel += 1
 195     end
 196     #p num
 197   end
 198   puts "deleted #{numdel} / #{num} messages"
 199 end
 200
 201 if optimize
 202   puts "optimizing index..."
 203   optt = time { index.index.optimize }
 204   puts "optimized index of size #{index.size} in #{optt}s."
 205 end