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     end
  86   end
  87
  88   unless username && password
  89     username = ask("Username for #{uri.host}: ");
  90     password = ask("Password for #{uri.host}: ") { |q| q.echo = false }
  91     puts # why?
  92   end
  93
  94   [username, password]
  95 end
  96
  97
  98 educate_user if ARGV.member? '--help'
  99
 100 archive = ARGV.delete "--archive"
 101 unusual = ARGV.delete "--unusual"
 102 force_archive = ARGV.delete "--force-archive"
 103 force_read = ARGV.delete "--force-read"
 104 the_usual = ARGV.delete "--the-usual"
 105 rebuild = ARGV.delete "--rebuild"
 106 force_rebuild = ARGV.delete "--force-rebuild"
 107 optimize = ARGV.delete "--optimize"
 108 start_at = # ok really need to use optparse or something now
 109   if(i = ARGV.index("--start-at"))
 110     raise "start-at requires a numeric argument: #{ARGV[i + 1].inspect}" unless ARGV.length > (i + 1) && ARGV[i + 1] =~ /\d/
 111     ARGV.delete_at i
 112     ARGV.delete_at(i).to_i # whoa!
 113   end
 114
 115 if(o = ARGV.find { |x| x =~ /^--/ })
 116   $stderr.puts "error: unknown option #{o}"
 117   educate_user
 118 end
 119
 120 $terminal.wrap_at = :auto
 121 Redwood::start
 122 index = Redwood::Index.new
 123 index.load
 124
 125 sources = ARGV.map do |uri|
 126   uri = "mbox://#{uri}" unless uri =~ %r!://!
 127   source = index.source_for uri
 128   unless source
 129     source =
 130       case uri
 131       when %r!^mbox\+ssh://!
 132         say "For SSH connections, if you will use public key authentication, you may leave the username and password blank."
 133         say "\n"
 134         username, password = get_login_info uri, index.sources
 135         Redwood::MBox::SSHLoader.new(uri, username, password, nil, !unusual, !!archive)
 136       when %r!^imaps?://!
 137         username, password = get_login_info uri, index.sources
 138         Redwood::IMAP.new(uri, username, password, nil, !unusual, !!archive)
 139       else
 140         Redwood::MBox::Loader.new(uri, nil, !unusual, !!archive)
 141       end
 142     index.add_source source
 143   end
 144   source
 145 end
 146
 147 sources = (sources + index.usual_sources).uniq if the_usual
 148 if rebuild || force_rebuild
 149   if start_at
 150     sources.each { |s| s.seek_to! start_at }
 151   else
 152     sources.each { |s| s.reset! }
 153   end
 154 end
 155
 156 found = {}
 157 start = Time.now
 158 begin
 159   sources.each do |source|
 160     if source.broken?
 161       $stderr.puts "error loading messages from #{source}: #{source.broken_msg}"
 162       next
 163     end
 164     next if source.done?
 165     puts "loading from #{source}... "
 166     num = 0
 167     start_offset = nil
 168     source.each do |offset, labels|
 169       start_offset ||= offset
 170       labels -= [:inbox] if force_archive
 171       labels -= [:unread] if force_read
 172       begin
 173         m = Redwood::Message.new :source => source, :source_info => offset, :labels => labels
 174         if found[m.id]
 175           puts "skipping duplicate message #{m.id}"
 176           next
 177         else
 178           found[m.id] = true
 179         end
 180
 181         m.remove_label :unread if m.status == "RO" unless force_read
 182         puts "# message at #{offset}, labels: #{labels * ', '}"
 183         if (rebuild || force_rebuild) &&
 184             (docid, entry = index.load_entry_for_id(m.id)) && entry
 185           if force_rebuild || entry[:source_info].to_i != offset
 186             puts "replacing message #{m.id} labels #{entry[:label].inspect} (offset #{entry[:source_info]} => #{offset})"
 187             m.labels = entry[:label].split.map { |l| l.intern }
 188             num += 1 if index.update_message m, source, offset
 189           end
 190         else
 191           num += 1 if index.add_message m
 192         end
 193       rescue Redwood::MessageFormatError, Redwood::SourceError => e
 194         $stderr.puts "ignoring erroneous message at #{source}##{offset}: #{e.message}"
 195       end
 196       if num % 1000 == 0 && num > 0
 197         elapsed = Time.now - start
 198         pctdone = (offset.to_f - start_offset) / (source.total.to_f - start_offset)
 199         remaining = (source.total.to_f - offset.to_f) * (elapsed.to_f / (offset.to_f - start_offset))
 200         puts "## #{num} (#{(pctdone * 100.0)}% done) read; #{elapsed.to_time_s} elapsed; est. #{remaining.to_time_s} remaining"
 201       end
 202     end
 203     puts "loaded #{num} messages" unless num == 0
 204   end
 205 ensure
 206   $stderr.puts "saving index and sources..."
 207   index.save
 208   Redwood::finish
 209 end
 210
 211 if rebuild || force_rebuild
 212   puts "deleting missing messages from the index..."
 213   numdel = num = 0
 214   sources.each do |source|
 215     raise "no source id for #{source}" unless source.id
 216     q = "+source_id:#{source.id}"
 217     q += " +source_info: >= #{start_at}" if start_at
 218     #p q
 219     num += index.index.search_each(q, :limit => :all) do |docid, score|
 220       mid = index.index[docid][:message_id]
 221       next if found[mid]
 222       puts "deleting #{mid}"
 223       index.index.delete docid
 224       numdel += 1
 225     end
 226     #p num
 227   end
 228   puts "deleted #{numdel} / #{num} messages"
 229 end
 230
 231 if optimize
 232   puts "optimizing index..."
 233   optt = time { index.index.optimize }
 234   puts "optimized index of size #{index.size} in #{optt}s."
 235 end