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