bin/mog

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