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

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