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