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