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

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