td 0.7.0

data/lib/td/command/account.rb

@@ -0,0 +1,84 @@
+
+module TD
+module Command
+
+  def account
+    op = cmd_opt 'account', :user_name?
+
+    op.banner << "\noptions:\n"
+
+    force = false
+    op.on('-f', '--force', 'overwrite current setting', TrueClass) {|b|
+      force = true
+    }
+
+    user_name = op.cmd_parse
+
+    require 'td/config'
+    conf = nil
+    begin
+      conf = Config.read($TRD_CONFIG_PATH)
+    rescue ConfigError
+    end
+    if conf && conf['account.user']
+      unless force
+        $stderr.puts "TreasureData account is already configured with '#{conf['account.user']}' account."
+        $stderr.puts "Add '-f' option to overwrite this setting."
+        exit 0
+      end
+    end
+
+    unless user_name
+      print "User name: "
+      line = STDIN.gets || ""
+      user_name = line.strip
+    end
+
+    if user_name.empty?
+      $stderr.puts "Canceled."
+      exit 0
+    end
+
+    api = nil
+
+    2.times do
+      begin
+        system "stty -echo"  # TODO termios
+        print "Password: "
+        password = STDIN.gets || ""
+        password = password[0..-2]  # strip \n
+      ensure
+        system "stty echo"   # TODO termios
+        print "\n"
+      end
+
+      if password.empty?
+        $stderr.puts "Canceled."
+        exit 0
+      end
+
+      require 'td/api'
+
+      begin
+        api = API.authenticate(user_name, password)
+      rescue TD::AuthError
+        $stderr.puts "User name or password mismatched."
+      end
+
+      break if api
+    end
+    return unless api
+
+    $stderr.puts "Authenticated successfully."
+
+    conf ||= Config.new
+    conf["account.user"] = user_name
+    conf["account.apikey"] = api.apikey
+    conf.save($TRD_CONFIG_PATH)
+
+    $stderr.puts "Use '#{$prog} create-database <db_name>' to create a database."
+  end
+
+end
+end
+

data/lib/td/command/common.rb

@@ -0,0 +1,121 @@
+
+module TD
+module Command
+  private
+  def cmd_opt(name, *args)
+    if args.last.to_s =~ /_$/
+      multi = true
+      args.push args.pop.to_s[0..-2]+'...'
+    end
+
+    req_args, opt_args = args.partition {|a| a.to_s !~ /\?$/ }
+    opt_args = opt_args.map {|a| a.to_s[0..-2].to_sym }
+    args = req_args + opt_args
+
+    args_line = req_args.map {|a| "<#{a}>" }
+    args_line.concat opt_args.map {|a| "[#{a}]" }
+    args_line = args_line.join(' ')
+
+    description = List.get_description(name)
+
+    op = OptionParser.new
+    op.summary_indent = "  "
+    op.banner = <<EOF
+usage: #{$prog} #{name} #{args_line}
+
+description:
+#{description.split("\n").map {|l| "  #{l}" }.join("\n")}
+EOF
+
+    (class<<op;self;end).module_eval do
+      define_method(:cmd_usage) do |msg|
+        puts op.to_s
+        puts ""
+        puts "error: #{msg}" if msg
+        exit 1
+      end
+
+      define_method(:cmd_parse) do
+        begin
+          parse!(ARGV)
+          if ARGV.length < req_args.length || (!multi && ARGV.length > args.length)
+            cmd_usage nil
+          end
+          if ARGV.length <= 1
+            ARGV[0]
+          else
+            ARGV
+          end
+        rescue
+          cmd_usage $!
+        end
+      end
+
+    end
+
+    op
+  end
+
+  def cmd_config
+    require 'td/config'
+    Config.read($TRD_CONFIG_PATH)
+  end
+
+  def cmd_api(conf)
+    apikey = conf['account.apikey']
+    unless apikey
+      raise ConfigError, "Account is not configured."
+    end
+    require 'td/api'
+    api = API.new(apikey)
+  end
+
+  def cmd_render_table(rows, *opts)
+    require 'hirb'
+    Hirb::Helpers::Table.render(rows, *opts)
+  end
+
+  def cmd_render_tree(nodes, *opts)
+    require 'hirb'
+    Hirb::Helpers::Tree.render(nodes, *opts)
+  end
+
+  def cmd_debug_error(ex)
+    if $verbose
+      $stderr.puts "error: #{$!.class}: #{$!.to_s}"
+      $!.backtrace.each {|b|
+        $stderr.puts "  #{b}"
+      }
+        $stderr.puts ""
+    end
+  end
+
+  def find_database(api, db_name)
+    begin
+      return api.database(db_name)
+    rescue
+      cmd_debug_error $!
+      $stderr.puts $!
+      $stderr.puts "Use '#{$prog} show-databases' to show the list of databases."
+      exit 1
+    end
+    db
+  end
+
+  def find_table(api, db_name, table_name, type=nil)
+    db = find_database(api, db_name)
+    begin
+      table = db.table(table_name)
+    rescue
+      $stderr.puts $!
+      $stderr.puts "Use '#{$prog} show-tables #{db_name}' to show the list of tables."
+      exit 1
+    end
+    if type && table.type != type
+      $stderr.puts "Table '#{db_name}.#{table_name} is not a #{type} table but a #{table.type} table"
+    end
+    table
+  end
+end
+end
+

data/lib/td/command/database.rb

@@ -0,0 +1,82 @@
+
+module TD
+module Command
+
+  def create_database
+    op = cmd_opt 'create-database', :db_name
+    db_name = op.cmd_parse
+
+    conf = cmd_config
+    api = cmd_api(conf)
+
+    begin
+      api.create_database(db_name)
+    rescue AlreadyExistsError
+      $stderr.puts "Database '#{db_name}' already exists."
+      exit 1
+    end
+
+    $stderr.puts "Database '#{db_name}' is created."
+    $stderr.puts "Use '#{$prog} create-log-table #{db_name} <table_name>' to create a table."
+  end
+
+  def drop_database
+    op = cmd_opt 'drop-database', :db_name
+
+    op.banner << "\noptions:\n"
+
+    force = false
+    op.on('-f', '--force', 'clear tables and delete the database', TrueClass) {|b|
+      force = true
+    }
+
+    db_name = op.cmd_parse
+
+    conf = cmd_config
+    api = cmd_api(conf)
+
+    begin
+      db = api.database(db_name)
+
+      if !force && !db.tables.empty?
+        $stderr.puts "Database '#{db_name}' is not empty. Use '-f' option or drop tables first."
+        exit 1
+      end
+
+      db.delete
+    rescue NotFoundError
+      $stderr.puts "Database '#{db_name}' does not exist."
+      eixt 1
+    end
+
+    $stderr.puts "Database '#{db_name}' is deleted."
+  end
+
+  def show_databases
+    op = cmd_opt 'show-databases'
+    op.cmd_parse
+
+    conf = cmd_config
+    api = cmd_api(conf)
+
+    dbs = api.databases
+
+    rows = []
+    dbs.each {|db|
+      rows << {:Name => db.name}
+    }
+    puts cmd_render_table(rows, :fields => [:Name])
+
+    if dbs.empty?
+      $stderr.puts "There are no databases."
+      $stderr.puts "Use '#{$prog} create-database <db_name>' to create a database."
+    end
+  end
+
+  alias show_dbs show_databases
+  alias create_db create_database
+  alias drop_db drop_database
+
+end
+end
+

data/lib/td/command/import.rb

@@ -0,0 +1,286 @@
+
+module TD
+module Command
+
+  IMPORT_TEMPLATES = {
+    'apache' => [
+                  /^([^ ]*) [^ ]* ([^ ]*) \[([^\]]*)\] "(\S+)(?: +([^ ]*) +\S*)?" ([^ ]*) ([^ ]*)(?: "([^\"]*)" "([^\"]*)")?$/,
+                  ['host', 'user', 'time', 'method', 'path', 'code', 'size', 'referer', 'agent'],
+                  "%d/%b/%Y:%H:%M:%S %z"],
+    'syslog' => [
+                  /^([^ ]* [^ ]* [^ ]*) ([^ ]*) ([a-zA-Z0-9_\/\.\-]*)(?:\[([0-9]+)\])?[^\:]*\: *(.*)$/,
+                  ['time', 'host', 'ident', 'pid', 'message'],
+                  "%b %d %H:%M:%S"],
+  }
+
+  # TODO import-item
+  # TODO tail
+
+  def import
+    op = cmd_opt 'import', :db_name, :table_name, :files_
+
+    op.banner << "\nsupported formats:\n"
+    op.banner << "  apache\n"
+    op.banner << "  syslog\n"
+    op.banner << "  msgpack\n"
+    op.banner << "  json\n"
+
+    op.banner << "\noptions:\n"
+
+    format = 'apache'
+    time_key = nil
+
+    op.on('--format FORMAT', "file format (default: #{format})") {|s|
+      format = s
+    }
+
+    op.on('--apache', "same as --format apache; apache common log format") {
+      format = 'apache'
+    }
+
+    op.on('--syslog', "same as --format syslog; syslog") {
+      format = 'syslog'
+    }
+
+    op.on('--msgpack', "same as --format msgpack; msgpack stream format") {
+      format = 'msgpack'
+    }
+
+    op.on('--json', "same as --format json; LF-separated json format") {
+      format = 'json'
+    }
+
+    op.on('-t', '--time-key COL_NAME', "time key name for json and msgpack format (e.g. 'created_at')") {|s|
+      time_key = s
+    }
+
+    db_name, table_name, *paths = op.cmd_parse
+
+    conf = cmd_config
+    api = cmd_api(conf)
+
+    case format
+    when 'json', 'msgpack'
+      unless time_key
+        $stderr.puts "-t, --time-key COL_NAME (e.g. '-t created_at') parameter is required for #{format} format"
+        exit 1
+      end
+      if format == 'json'
+        require 'json'
+        require 'time'
+        parser = JsonParser.new(time_key)
+      else
+        parser = MessagePackParser.new(time_key)
+      end
+
+    else
+      regexp, names, time_format = IMPORT_TEMPLATES[format]
+      if !regexp || !names || !time_format
+        $stderr.puts "Unknown format '#{format}'"
+        exit 1
+      end
+      parser = TextParser.new(names, regexp, time_format)
+    end
+
+    find_table(api, db_name, table_name, :log)
+
+    require 'zlib'
+
+    files = paths.map {|path|
+      if path == '-'
+        $stdin
+      elsif path =~ /\.gz$/
+        Zlib::GzipReader.open(path)
+      else
+        File.open(path)
+      end
+    }
+
+    require 'msgpack'
+    require 'tempfile'
+    #require 'thread'
+
+    files.zip(paths).each {|file,path|
+      import_log_file(file, path, api, db_name, table_name, parser)
+    }
+
+    puts "done."
+  end
+
+  private
+  def import_log_file(file, path, api, db_name, table_name, parser)
+    puts "importing #{path}..."
+
+    out = Tempfile.new('td-import')
+    writer = Zlib::GzipWriter.new(out)
+
+    n = 0
+    x = 0
+    parser.call(file, path) {|record|
+      writer.write record.to_msgpack
+
+      n += 1
+      x += 1
+      if n % 10000 == 0
+        puts "  imported #{n} entries from #{path}..."
+
+      elsif out.pos > 1024*1024  # TODO size
+        puts "  imported #{n} entries from #{path}..."
+        begin
+          writer.finish
+          size = out.pos
+          out.pos = 0
+
+          puts "  uploading #{size} bytes..."
+          api.import(db_name, table_name, "msgpack.gz", out, size)
+
+          out.truncate(0)
+          out.pos = 0
+          x = 0
+          writer = Zlib::GzipWriter.new(out)
+        rescue
+          $stderr.puts "  #{$!}"
+          return 1 # TODO error
+        end
+      end
+    }
+
+    if x != 0
+      writer.finish
+      size = out.pos
+      out.pos = 0
+
+      puts "  uploading #{size} bytes..."
+      # TODO upload on background thread
+      api.import(db_name, table_name, "msgpack.gz", out, size)
+    end
+
+    puts "  imported #{n} entries from #{path}."
+
+  ensure
+    out.close rescue nil
+    writer.close rescue nil
+  end
+
+  require 'date'  # DateTime#strptime
+  require 'time'  # Time#strptime, Time#parse
+
+  class TextParser
+    def initialize(names, regexp, time_format)
+      @names = names
+      @regexp = regexp
+      @time_format = time_format
+    end
+
+    def call(file, path, &block)
+      i = 0
+      file.each_line {|line|
+        i += 1
+        begin
+          m = @regexp.match(line)
+          unless m
+            raise "invalid log format at #{path}:#{i}"
+          end
+
+          record = {}
+
+          cap = m.captures
+          @names.each_with_index {|name,i|
+            if value = cap[i]
+              if name == "time"
+                value = parse_time(value).to_i
+              end
+              record[name] = value
+            end
+          }
+
+          block.call(record)
+
+        rescue
+          $stderr.puts "  skipped: #{$!}: #{line.dump}"
+        end
+      }
+    end
+
+    if Time.respond_to?(:strptime)
+      def parse_time(value)
+        Time.strptime(value, @time_format)
+      end
+    else
+      def parse_time(value)
+        Time.parse(DateTime.strptime(value, @time_format).to_s)
+      end
+    end
+  end
+
+  class JsonParser
+    def initialize(time_key)
+      require 'json'
+      @time_key = time_key
+    end
+
+    def call(file, path, &block)
+      i = 0
+      file.each_line {|line|
+        i += 1
+        begin
+          record = JSON.parse(line)
+
+          time = record[@time_key]
+          unless time
+            raise "record doesn't have '#{@time_key}' column"
+          end
+
+          case time
+          when Integer
+            # do nothing
+          else
+            time = Time.parse(time.to_s).to_i
+          end
+          record['time'] = time
+
+          block.call(record)
+
+        rescue
+          $stderr.puts "  skipped: #{$!}: #{line.dump}"
+        end
+      }
+    end
+  end
+
+  class MessagePackParser
+    def initialize(time_key)
+      require 'json'
+      @time_key = time_key
+    end
+
+    def call(file, path, &block)
+      i = 0
+      MessagePack::Unpacker.new(file).each {|record|
+        i += 1
+        begin
+          time = record[@time_key]
+          unless time
+            raise "record doesn't have '#{@time_key}' column"
+          end
+
+          case time
+          when Integer
+            # do nothing
+          else
+            time = Time.parse(time.to_s).to_i
+          end
+          record['time'] = time
+
+          block.call(record)
+
+        rescue
+          $stderr.puts "  skipped: #{$!}: #{record.to_json}"
+        end
+      }
+    rescue EOFError
+    end
+  end
+end
+end
+