bin/mog

   1 #!/usr/bin/env ruby
   2 require 'mogilefs'
   3 require 'optparse'
   4 $stdin.binmode
   5 $stdout.binmode
   6 $stderr.sync = $stdout.sync = true
   7
   8 trap('INT') { exit 130 }
   9 trap('PIPE') { exit 0 }
  10
  11 # this is to be compatible with config files used by the Perl tools
  12 def parse_config_file!(path, overwrite = false)
  13   dest = {}
  14   File.open(path).each_line do |line|
  15     line.strip!
  16     if /^(domain|class)\s*=\s*(\S+)/.match(line)
  17       dest[$1.to_sym] = $2
  18     elsif m = /^(?:trackers|hosts)\s*=\s*(.*)/.match(line)
  19       dest[:hosts] = $1.split(/\s*,\s*/)
  20     elsif m = /^timeout\s*=\s*(.*)/.match(line)
  21       dest[:timeout] = m[1].to_f
  22     else
  23       warn "Ignored configuration line: #{line}" unless /^#/.match(line)
  24     end
  25   end
  26   dest
  27 end
  28
  29 # parse the default config file if one exists
  30 def_file = File.expand_path("~/.mogilefs-client.conf")
  31 def_cfg = File.exist?(def_file) ? parse_config_file!(def_file) : {}
  32
  33 # parse the command-line first, these options take precedence over all else
  34 cli_cfg = {}
  35 config_file = nil
  36 ls_l = false
  37 ls_h = false
  38 chunk = false
  39 test = {}
  40 cat = { :raw => false }
  41
  42 ARGV.options do |x|
  43   x.banner = "Usage: #{$0} [options] <command> [<arguments>]"
  44   x.separator ''
  45
  46   x.on('-c', '--config=/path/to/config',
  47        'config file to load') { |file| config_file = file }
  48
  49   x.on('-t', '--trackers=host1[,host2]', '--hosts=host1[,host2]', Array,
  50        'hostnames/IP addresses of trackers') do |trackers|
  51     cli_cfg[:hosts] = trackers
  52   end
  53
  54   x.on('-e', 'True if key exists') { test[:e] = true }
  55   x.on('-r', '--raw', 'show raw big_info file information') { cat[:raw] = true }
  56
  57   x.on('-C', '--class=s', 'class') { |klass| cli_cfg[:class] = klass }
  58   x.on('-d', '--domain=s', 'domain') { |domain| cli_cfg[:domain] = domain }
  59   x.on('-l', "long listing format (`ls' command)") { ls_l = true }
  60   x.on('-h', '--human-readable',
  61        "print sizes in human-readable format (`ls' command)") { ls_h = true }
  62   x.on('--chunk', "chunk uploads (`tee' command)") { chunk = true }
  63   x.separator ''
  64   x.on('--help', 'Show this help message.') { puts x; exit }
  65   x.on('--version', 'Show --version') { puts "#$0 #{MogileFS::VERSION}"; exit }
  66   x.parse!
  67 end
  68
  69 # parse the config file specified at the command-line
  70 file_cfg = config_file ? parse_config_file!(config_file, true) : {}
  71
  72 # read environment variables, too.  This Ruby API favors the term
  73 # "hosts", however upstream MogileFS teminology favors "trackers" instead.
  74 # Favor the term more consistent with what the MogileFS inventors used.
  75 env_cfg = {}
  76 if ENV["MOG_TRACKERS"]
  77   env_cfg[:hosts] = ENV["MOG_TRACKERS"].split(/\s*,\s*/)
  78 end
  79 if ENV["MOG_HOSTS"] && (env_cfg[:hosts] || []).empty?
  80   env_cfg[:hosts] = ENV["MOG_HOSTS"].split(/\s*,\s*/)
  81 end
  82 env_cfg[:domain] = ENV["MOG_DOMAIN"] if ENV["MOG_DOMAIN"]
  83 env_cfg[:class] = ENV["MOG_CLASS"] if ENV["MOG_CLASS"]
  84
  85 # merge the configs, favoring them in order specified:
  86 cfg = {}.merge(def_cfg).merge(env_cfg).merge(file_cfg).merge(cli_cfg)
  87
  88 # error-checking
  89 err = []
  90 err << "trackers must be specified" if cfg[:hosts].nil? || cfg[:hosts].empty?
  91 err << "domain must be specified" unless cfg[:domain]
  92 if err.any?
  93   warn "Errors:\n  #{err.join("\n  ")}"
  94   warn ARGV.options
  95   exit 1
  96 end
  97
  98 unless cmd = ARGV.shift
  99   warn ARGV.options
 100   exit 1
 101 end
 102
 103 cfg[:timeout] ||= 30 # longer timeout for interactive use
 104 mg = MogileFS::MogileFS.new(cfg)
 105
 106 def store_file_retry(mg, key, storage_class, filepath)
 107   tries = 0
 108   begin
 109     mg.store_file(key, storage_class, filepath)
 110   rescue MogileFS::UnreadableSocketError,
 111          MogileFS::Backend::NoDevicesError => err
 112     if ((tries += 1) < 10)
 113       warn "Retrying on error: #{err}: #{err.message} tries: #{tries}"
 114       retry
 115     else
 116       warn "FATAL: #{err}: #{err.message} tries: #{tries}"
 117     end
 118     exit 1
 119   end
 120 end
 121
 122 def human_size(size)
 123   suff = ''
 124   %w(K M G).each do |s|
 125     size /= 1024.0
 126     if size <= 1024
 127       suff = s
 128       break
 129     end
 130   end
 131   sprintf("%.1f%s", size, suff)
 132 end
 133
 134 begin
 135   case cmd
 136   when 'cp'
 137     filename = ARGV.shift or raise ArgumentError, '<filename> <key>'
 138     dkey = ARGV.shift or raise ArgumentError, '<filename> <key>'
 139     ARGV.shift and raise ArgumentError, '<filename> <key>'
 140     store_file_retry(mg, dkey, cfg[:class], filename)
 141   when 'cat'
 142     ARGV.empty? and raise ArgumentError, '<key1> [<key2> ...]'
 143     ARGV.each do |key|
 144       if (!cat[:raw] && key =~ /^_big_info:/)
 145         mg.bigfile_write(key, $stdout, {:verify => true})
 146       else
 147         mg.get_file_data(key, $stdout)
 148       end
 149     end
 150   when 'ls'
 151     prefixes = ARGV.empty? ? [ nil ] : ARGV
 152     if ls_l
 153       each_key = lambda do |key, size, devcount|
 154         size = ls_h && size > 1024 ? human_size(size) : size.to_s
 155         size = (' ' * (12 - size.length)) << size # right justify
 156         puts [ sprintf("% 2d", devcount), size, key ].pack("A4 A16 A*")
 157       end
 158     else
 159       each_key = lambda { |key| puts key }
 160     end
 161     prefixes.each { |prefix| mg.each_key(prefix, &each_key) }
 162   when 'rm'
 163     ARGV.empty? and raise ArgumentError, '<key1> [<key2>]'
 164     ARGV.each { |key| mg.delete(key) }
 165   when 'mv'
 166     from = ARGV.shift or raise ArgumentError, '<from> <to>'
 167     to = ARGV.shift or raise ArgumentError, '<from> <to>'
 168     ARGV.shift and raise ArgumentError, '<from> <to>'
 169     mg.rename(from, to)
 170   when 'stat' # this outputs a RFC822-like format
 171     ARGV.empty? and raise ArgumentError, '<key1> [<key2>]'
 172     ok = true
 173     ARGV.each_with_index do |key,j|
 174       begin
 175         info = mg.file_info(key)
 176         puts "Key: #{key}"
 177         puts "Size: #{info['length']}"
 178         puts "Class: #{info['class']}"
 179         o = { :pathcount => info["devcount"] }
 180         mg.get_paths(key, o).each_with_index do |path,i|
 181           puts "URL-#{i}: #{path}"
 182         end
 183         puts "" if ARGV.size != (j + 1)
 184       rescue MogileFS::Backend::UnknownKeyError
 185         warn "No such key: #{key}"
 186         ok = false
 187       end
 188     end
 189     exit(ok)
 190   when 'tee'
 191     require 'tempfile'
 192     dkey = ARGV.shift or raise ArgumentError, '<key>'
 193     ARGV.shift and raise ArgumentError, '<key>'
 194     skip_tee = File.stat('/dev/null') == $stdout.stat
 195
 196     if chunk
 197       if skip_tee
 198         tee_obj = $stdin
 199       else
 200         tee_obj = lambda do |*args|
 201           buf = $stdin.readpartial(*args)
 202           $stdout.write(buf)
 203           buf
 204         end
 205         class << tee_obj
 206           alias readpartial call
 207         end
 208       end
 209       mg.store_file(dkey, cfg[:class], tee_obj)
 210     else # buffer input, first
 211       tmp = Tempfile.new('mog-tee')
 212       tmp.sync = true
 213
 214       # if stdout is pointing to /dev/null, don't bother installing the filter.
 215       tee_obj = tmp
 216       unless skip_tee
 217         tee_obj = lambda do |buf|
 218           $stdout.write(buf)
 219           tmp.write(buf)
 220         end
 221         class << tee_obj
 222           alias write call
 223         end
 224       end
 225       begin
 226         MogileFS.io.copy_stream($stdin, tee_obj)
 227         store_file_retry(mg, dkey, cfg[:class], tmp.path)
 228       ensure
 229         tmp.close!
 230       end
 231     end
 232   when 'test'
 233     truth, ok = true, nil
 234     raise ArgumentError, "-e must be specified" unless (test.size == 1)
 235
 236     truth, key = case ARGV.size
 237     when 1
 238       [ true, ARGV[0] ]
 239     when 2
 240       if ARGV[0] != "!"
 241         raise ArgumentError, "#{ARGV[0]}: binary operator expected"
 242       end
 243       [ false, ARGV[1] ]
 244     else
 245       raise ArgumentError, "Too many arguments"
 246     end
 247
 248     begin
 249       paths = mg.get_paths(key)
 250       if test[:e]
 251         ok = !!(paths && paths.size > 0)
 252       else
 253         raise ArgumentError, "Unknown flag: -#{test.keys.first}"
 254       end
 255     rescue MogileFS::Backend::UnknownKeyError
 256       ok = false
 257     end
 258
 259     truth or ok = ! ok
 260     exit ok ? 0 : 1
 261   else
 262     raise ArgumentError, "Unknown command: #{cmd}"
 263   end
 264 rescue ArgumentError => err
 265   warn "Usage: #{$0} #{cmd} #{err.message}"
 266   exit 1
 267 end
 268 exit 0