iij-dag 1.0.1

data/lib/dag/cli/command/import.rb

@@ -0,0 +1,30 @@
+module Dag
+  module CLI
+    class Import
+      def initialize(client, db, table, files, format, params)
+        @client = client
+        @db = db
+        @table = table
+        @files = files
+        @params = params
+        @format = format
+      end
+
+      def execute
+        begin
+          @client.database(@db)
+        rescue Dag::Client::DatabaseNotFound
+          @client.create_database(@db)
+        end
+
+        begin
+          @client.database(@db).table(@table)
+        rescue Dag::Client::TableNotFound
+          @client.create_table(@db, @table, format: @format)
+        end
+
+        @client.import(@db, @table, @files, @params)
+      end
+    end
+  end
+end

data/lib/dag/cli/command/query.rb

@@ -0,0 +1,14 @@
+module Dag
+  module CLI
+    class Query
+      def initialize(client, params)
+        @client = client
+        @params = params
+      end
+
+      def execute
+        @client.query(@params)
+      end
+    end
+  end
+end

data/lib/dag/cli/sub_command.rb

@@ -0,0 +1,18 @@
+module Dag
+  module CLI
+    class SubCommand < Base
+      include Dag::CLI::Utils::TimeFormat
+
+      def self.banner(task, namespace = false, subcommand = true)
+        super
+      end
+
+      desc 'help [COMMAND]', 'Describe available commands or one specific command'
+      def help(command = nil, subcommand = false)
+        super
+        cmd = self.class.name.split('::').last.downcase
+        shell.say "Type 'dag #{cmd} help [COMMAND]' for more information on a specific command"
+      end
+    end
+  end
+end

data/lib/dag/cli/sub_command/bucket.rb

@@ -0,0 +1,90 @@
+# coding: utf-8
+
+require 'pathname'
+
+module Dag
+  module CLI
+    class Bucket < SubCommand
+      desc 'list ([BUCKET])', 'show bucket/object list'
+      option :prefix, type: :string, aliases: '-p', desc: "prefix filter", default: nil
+      option :delimiter, type: :string, aliases: '-d', desc: "object delimiter", default: nil
+      def list(bucket = nil)
+        fields = [:name]
+
+        unless bucket
+          rows = handle_api_failure do
+            client.buckets
+          end
+        else
+          rows = handle_api_failure do
+            bucket = client.buckets[bucket]
+            bucket.objects.where(prefix: options[:prefix], delimiter: options[:delimiter])
+          end
+        end
+
+        handle_api_failure do
+          rows.each do |row|
+            puts row.name
+          end
+        end
+      end
+
+      desc 'read [BUCKET] [OBJECT]', 'read object'
+      def read(bucket, object)
+        rows = handle_api_failure do
+          bucket = client.buckets[bucket]
+          bucket.objects[object].read
+        end
+        print rows
+      end
+
+      desc 'create [BUCKET]', 'create bucket'
+      def create(bucket)
+        rows = handle_api_failure do
+          client.create_bucket(bucket)
+        end
+        say_status("Create bucket", rows.name)
+      end
+
+      desc 'delete [BUCKET] ([OBJECT])', 'delete bucket or object'
+      def delete(bucket, *objects)
+        handle_api_failure do
+          if objects.empty?
+            # for bucket
+            client.delete_bucket(bucket)
+            say_status("Delete bucket", bucket)
+          else
+            # for objects
+            th = []
+            objects.each do |object|
+              th << Thread.new do
+                client.delete_object(bucket, object)
+                say_status("Delete object", File.join(bucket, object))
+              end
+            end
+            th.map(&:join)
+          end
+        end
+      end
+
+      desc 'write [BUCKET] [OBJECT]', 'write object / string or file path'
+      option :data, type: :string, aliases: '-d', desc: "string or file path", required: true
+      option :multipart, type: :boolean, aliases: '-m', desc: "multipart upload"
+      option :multipart_jobs, type: :numeric, aliases: '-j', desc: "multipart upload jobs"
+      option :multipart_splitsz, type: :numeric, aliases: '-s', desc: "multipart upload split size"
+      def write(bucket, object)
+        params = {}
+        params.merge!(multipart: true) if options[:multipart]
+        params.merge!(jobs: options[:multipart_jobs]) if options[:multipart_jobs]
+        params.merge!(splitsz: options[:multipart_splitsz]) if options[:multipart_splitsz]
+
+        data = File.exist?(options[:data]) ? Pathname.new(options[:data]) : options[:data]
+        handle_api_failure do
+          obj = client.buckets[bucket].objects[object]
+          obj.write(data, params)
+          say_status("Write object", File.join(bucket, object))
+        end
+      end
+    end
+  end
+end

data/lib/dag/cli/sub_command/cluster.rb

@@ -0,0 +1,187 @@
+# coding: utf-8
+module Dag
+  module CLI
+    class Cluster < SubCommand
+      include Dag::CLI::Utils::NumberToHuman
+
+      class NotClusterFound < StandardError; end
+
+      desc 'attach ([CLUSTER])', 'Attach cluster'
+      def attach(cluster = nil)
+        if cluster
+          attached_cluster = handle_api_failure do
+            open_client{ cluster }.cluster
+          end
+
+          local_store { |db| db[:cluster] = cluster }
+          say_status "Attach cluster", attached_cluster.name
+        else
+          cluster = local_store { |db| db[:cluster] }
+          begin
+            raise NotClusterFound if cluster.nil?
+            say_status "Attached cluster", cluster
+          rescue NotClusterFound
+            abort "Please attach to the cluster\nType 'dag cluster attach [CLUSTER]'"
+          end
+        end
+      end
+
+      desc 'detach', 'Detach cluster'
+      def detach
+        attached_cluster = ''
+        local_store do |db|
+          attached_cluster = db[:cluster]
+          raise NotClusterFound if attached_cluster.nil?
+          db[:cluster] = nil
+        end
+        say_status "Detach cluster", attached_cluster
+      rescue NotClusterFound
+        abort "Please attach to the cluster\nType 'dag attach [CLUSTER]'"
+      end
+
+      desc 'list', 'show cluster list'
+      def list
+        headers = %w(# name status type instances)
+        rows = []
+        attach_cluster = local_store { |db| db[:cluster] }
+
+        handle_api_failure do
+          client.clusters.each do |cluster|
+            total_container = cluster.instances.lazy.map{|c| c["quantity"] }.inject(:+)
+            cluster_row = [cluster.name, cluster.status, cluster.type, total_container || 0]
+            rows << if attach_cluster == cluster.name
+              cluster_row.unshift('*')
+            else
+              cluster_row.unshift(nil)
+            end
+          end
+        end
+
+        terminal_table(rows, headers: headers, max_width: Float::INFINITY)
+      end
+
+      desc 'info', 'Cluster instance info id/dfs_used/non_dfs_used/capacity/grade'
+      def info
+        headers = %w(instance_id grade dfs_used non_dfs_used capacity)
+        rows = []
+
+        cluster = handle_api_failure do
+          open_client.cluster
+        end
+        statistics = cluster.statistics
+
+        unless statistics
+          return say_status(
+              "InvalidClusterState",
+              "Cluster status is invalid: #{cluster.status}",
+              :red
+          )
+        end
+
+        statistics.instances.each do |instance|
+          rows << [
+              instance.instance_id,
+              instance.grade,
+              number_to_human(instance.disk.dfs_used),
+              number_to_human(instance.disk.non_dfs_used),
+              number_to_human(instance.disk.capacity)
+          ]
+        end
+
+        puts "Cluster: #{cluster.name}"
+        terminal_table(rows, headers: headers) 
+      end
+
+      desc 'restart', 'restart cluster'
+      option :force, type: :boolean, aliases: '-f', default: false, desc: 'cluster force restart'
+      option :type,  type: :string,  aliases: '-t', default: nil, desc: 'change cluster type'
+      option :wait,  type: :boolean, aliases: '-w', default: false, desc: 'wait until normal cluster status'
+      option :assumeyes,   type: :boolean, aliases: '-y', desc: 'assume that the answer to any question which would be asked is yes'
+      def restart
+        unless options[:assumeyes]
+          loop do
+            ans = ask 'Is this ok [y/N]:'
+            case ans
+              when 'y'
+                break
+              when 'N'
+                exit
+            end
+          end
+        end
+
+        headers = %w(name type force)
+        rows = []
+
+        cluster = handle_api_failure do
+          open_client.cluster
+        end
+
+        options[:type] ||= cluster.type
+        params = { force: options[:force], type: options[:type] }
+
+        handle_api_failure do
+          cluster.restart(params)
+        end
+
+        rows << [cluster.name, params[:type], params[:force]]
+
+        terminal_table(rows, headers: headers)
+
+        if options[:wait]
+          puts
+          puts 'wait cluster ready...'
+          wait_ready
+        end
+      end
+
+      desc 'wait_ready', 'wait cluster ready'
+      def wait_ready
+        pretime = Time.now
+        prestate = ""
+        state = []
+        state_num = 0
+        while true do
+          cluster = handle_api_failure do
+            open_client.cluster
+          end
+
+          st = cluster.status
+          now = Proc.new { "#{st}(#{(Time.now-pretime).round}s)" }
+          fiber = Fiber.new {|state| Fiber.yield state }
+
+          if prestate != st
+            prestate = st
+            pretime = Time.now
+            state << now.call
+            state_num += 1
+          end
+
+          state[state_num-1] = now.call
+          print "#{fiber.resume(state*' ')}\r"
+          $stdout.flush
+
+          if %w(norm failed ptfailed error).include? st
+            print "\n"
+            break
+          end
+
+          sleep 1
+        end
+      end
+
+      desc 'log [OUTPUT_PATH]', 'export cluster log'
+      option :compress, type: :boolean, aliases: '-c', default: false
+      def log(output_log_path)
+        params = {
+            compress: options[:compress],
+            output_log_path: output_log_path
+        }
+        handle_api_failure do
+          cluster = open_client.cluster
+          cluster.export_log(params)
+        end
+      end
+    end
+  end
+end

data/lib/dag/cli/sub_command/db.rb

@@ -0,0 +1,38 @@
+# coding: utf-8
+module Dag
+  module CLI
+    class Db < SubCommand
+
+      desc 'list', 'Show database list'
+      def list
+        fields = %w(name)
+        rows = handle_api_failure { open_client.databases }
+        terminal_table(rows, fields: fields)
+      end
+
+      desc 'create [DATABASE]', 'Create database'
+      def create(db_name)
+        database = handle_api_failure do
+          open_client.create_database(db_name)
+        end
+
+        say_status "Create Database", database.name
+      end
+
+      desc 'delete [DATABASE]', 'Delete database'
+      def delete(db_name)
+        db = ''
+        handle_api_failure do
+          db = open_client.database(db_name)
+          db.delete if db
+        end
+
+        if db
+          say_status "Delete Database", db_name
+        else
+          say_status "Unknown Database", db_name, :red
+        end
+      end
+    end
+  end
+end

data/lib/dag/cli/sub_command/job.rb

@@ -0,0 +1,341 @@
+# coding: utf-8
+
+require 'ruby-progressbar'
+require 'csv'
+require 'anbt-sql-formatter/formatter'
+
+module Dag
+  module CLI
+    class Job < SubCommand
+
+      desc 'list', 'Jobs list'
+      option :order, type: :string, aliases: '-o', desc: "list order [asc|desc]"
+      option :filter, type: :hash, aliases: '-f', desc: "list filter key [status|type|cluster_name|label|cluster_rebooted]"
+      option :limit, type: :numeric, aliases: '-n', desc: "The location is number lines"
+      def list
+        headers = %w(id status start_at type cluster query)
+        rows = []
+
+        if options[:filter] && cluster_rebooted = options[:filter]['cluster_rebooted']
+          options[:filter]['cluster_rebooted'] = case cluster_rebooted
+            when "true"
+              true
+            when "false"
+              false
+          end
+        end
+
+        handle_api_failure do
+          jobs = client.jobs.where(cluster_name: open_client.cluster_name)
+
+          if options[:filter]
+            jobs = jobs.where(options[:filter].symbolize_keys)
+          end
+
+          if options[:limit]
+            jobs = jobs.limit(options[:limit])
+          end
+
+          iterator = :each
+          order = if options[:order]
+                    options[:order]
+                  else
+                    if [:limit]
+                      iterator = :reverse_each
+                      'desc'
+                    else
+                      'asc'
+                    end
+                  end
+
+          jobs.order(order).send(iterator) do |job|
+            row = [ job.id ]
+            row += [ job.status ]
+            row += [ time_format(job.start_at) ]
+            row += [ job.type ]
+            row += [ job.cluster ]
+            row += [ job.query ]
+            rows << row
+          end
+        end
+
+        terminal_table(rows, headers: headers)
+      end
+
+      desc 'info [JOB_ID]', 'Job info'
+      def info(job_id)
+        headers = ["item", "info"]
+        rows = []
+
+        job = handle_api_failure do
+          client.job(job_id)
+        end
+
+        rows << ["id", job.id]
+        rows << ["status", job.status]
+        rows << ["start_at", time_format(job.start_at)]
+        rows << ["type", job.type]
+        rows << ["dsl", job.dsl]
+        rows << ["cluster", job.cluster]
+        rows << ["cluster_rebooted", job.cluster_rebooted]
+        rows << ["label", job.label]
+        rows << ["stage", job.stage]
+        rows << ["progress", job.progress]
+        rows << ["access_key_id", job.access_key_id]
+
+        case job.dsl
+          when "hive"
+            rows << ["output_format", job.output_format]
+            rows << ["output_resource_path", job.output_resource_path]
+            query = sql_format(job.query).split("\n")
+            query.each_with_index do |line, i|
+              rows << ['', line]
+              rows.last[0] = 'query' if i == 0
+            end
+          when "mapreduce"
+            rows << ["job_id", job.job_id]
+            rows << ["schema", job.schema]
+            rows << ["input_object_keys", job.input_object_keys]
+            rows << ["input_format", job.input_format]
+            rows << ["output_database", job.output_database]
+            rows << ["output_table", job.output_table]
+        end
+
+        # vertical, unicode, tab, markdown, simple
+        terminal_table(rows, headers: headers, max_width: Float::INFINITY)
+      end
+
+      desc 'reuse [JOB_ID]', 'reuse job'
+      option :output, type: :string, aliases: '-o', desc: 'new output storage path'
+      option :format, type: :string, aliases: '-f', desc: 'new output format [csv/tsv]'
+      option :label, type: :string, aliases: '-l', desc: 'new job label'
+      option :query, type: :string, aliases: '-q', desc: 'new query'
+      option :wait, type: :boolean, aliases: '-w', desc: 'wait job finish'
+      def reuse(job_id)
+        reuse_job = handle_api_failure do
+          client.job(job_id)
+        end
+
+        if reuse_job.type == "select"
+          query_string = options[:query]
+          if query_string
+            if File.exists? query_string
+              query_string = open(query_string).read
+              puts "Query: #{query_string}"
+            end
+          else
+            query_string = reuse_job.query
+          end
+
+          output_format = options[:format]
+          unless output_format
+            output_format = reuse_job.output_format
+          end
+
+          output_resource_path = options[:output]
+          unless output_resource_path
+            output_resource_path = reuse_job.output_resource_path
+          end
+
+          label = options[:label]
+          unless label
+            label = reuse_job.label
+          end
+
+          params = {
+              query: query_string,
+              output_format: output_format,
+              output_resource_path: output_resource_path,
+              label: label
+          }
+
+          query = Query.new(open_client, params)
+          job = handle_api_failure { query.execute }
+
+          say_status "accepted job", "job_id: #{job.id}"
+          if options[:wait]
+            CLI::Command.start(["job", "log", job.id.to_s, "-t"])
+            CLI::Command.start(["job", "result", job.id.to_s])
+          end
+        else
+          say_status "invalid job type", "job_type: #{reuse_job.type}"
+          exit 1
+        end
+      end
+
+      desc 'log [JOB_ID]', 'show job log'
+      option :tail, type: :boolean, aliases: '-t', desc: 'tail -f'
+      def log(job_id)
+        if options[:tail]
+          mark=nil
+          while true
+            jst, log = handle_api_failure {
+              j = client.job(job_id)
+              [j.running?, j.log]
+            }
+            lines = log.split("\n")
+            if mark
+              n = lines.rindex(mark)
+              lines = lines[(n+1)..-1] if n
+            end
+            unless lines.empty?
+              mark=lines[-1]
+              puts lines.join("\n")
+            end
+            break unless jst
+            sleep 1
+          end
+        else
+          log = handle_api_failure do
+            client.job(job_id).log
+          end
+
+          puts log
+        end
+      end
+
+      desc 'result [JOB_ID]', 'Create download url'
+      option :format, type: :string, aliases: '-f', desc: 'output format [csv|tsv|table|url|time]', default: "table"
+      def result(job_id)
+        if options[:format] == "url"
+          download_url = handle_api_failure do
+            client.job(job_id).download_urls
+          end
+          puts download_url
+          return
+        end
+
+        unless job_id.include? "/"
+          job = handle_api_failure do
+            client.job(job_id)
+          end
+          target = job.output_resource_path
+        else
+          target = job_id
+        end
+
+        unless target
+          if job.try(:type) == 'split'
+            say_status 'ERROR', 'Invalid job type', :red
+          else
+            say_status 'ERROR', 'Invalid job_id or dag://bucket/prefix/', :red
+          end
+          exit 1
+        end
+
+        unless job && job.finished?
+          say_status 'ERROR', "Job is not finished (status=#{job.status})", :red
+          return
+        end
+
+        uri = URI.parse(target)
+        bucket = uri.host
+        prefix = uri.path[1..-1]
+        bucket = client.buckets[bucket]
+        rsm = handle_api_failure do
+          bucket.objects[File.join(prefix, ".resultsetmetadata")].read
+        end
+
+        if rsm.headers['Last-Modified'].present?
+          finishedat = Time.parse(rsm.headers['Last-Modified'].first)
+        end
+
+        if %w(table time).include? options[:format]
+          h = []
+          h << "job start at #{time_format(job.start_at)}" if job
+          h << "finished #{time_format(finishedat)}" if finishedat
+          h << "#{(finishedat-job.start_at).to_i} seconds." if job && finishedat
+          puts h.join(", ")
+          return if options[:format]=="time"
+        end
+
+        field_separator = ","
+        field_separator = "\t" if job.output_format == "tsv"
+
+        hdr = Oj.load(rsm)
+        hdrs=[]
+        hdr["columns"].each{|e| hdrs[e["position"].to_i-1] = e["name"] }
+        rows = handle_api_failure do
+          bucket.objects.where(prefix: prefix)
+           .select{|s| s.name.include? "/0" }
+           .map{|e| CSV.parse(bucket.objects[e.name].read, col_sep: field_separator) }
+        end
+
+        rows.each{|obj|
+          case options[:format]
+          when 'csv'
+            puts hdrs.join(",")
+            puts obj.map{|m| m.join(",") }.join("\n")
+          when 'tsv'
+            puts hdrs.join("\t")
+            puts obj.map{|m| m.join("\t") }.join("\n")
+          else
+            terminal_table(obj, headers: hdrs)
+          end
+        }
+      end
+
+      desc 'kill! [JOB_ID]', 'Stop job'
+      def kill!(job_id)
+        handle_api_failure do
+          client.job(job_id).kill
+        end
+
+        say_status("Cancelling job", job_id)
+      end
+
+      desc 'progress [JOB_ID]', 'Show Progress Bar'
+      def progress(job_id=nil)
+        if job_id
+          jobs = [job_id]
+        else
+          jobs = handle_api_failure{
+            client.jobs.select{|j| j.running? }.map{|j| j.id }
+          }
+        end
+        return unless jobs.length==1
+        progress_bar = ProgressBar.create(total: 100, format: '%t %a <%B> %p%% %e', autofinish: false)
+        stage = nil
+        while true
+          job = handle_api_failure{ client.job(jobs[0]) }
+          if ! job.running?
+            progress_bar.finish
+            puts "#{job.status}"
+            break
+          end
+          if job.stage && stage!=job.stage
+            progress_bar.finish if stage
+            progress_bar = ProgressBar.create(title: "Stage-#{job.stage}/Job-#{jobs[0]}", total: 100, format: '%t %a <%B> %p%% %e', autofinish: false)
+            stage = job.stage
+          end
+          progress_bar.progress = job.progress.to_f if job.progress && stage == job.stage
+          sleep 1
+        end
+      end
+
+      desc 'wait [JOB_ID]', 'wait job'
+      def wait(job_id=nil)
+        while true
+          if job_id
+            job = handle_api_failure{ client.job(job_id) }
+            break unless job.running?
+          else
+            jobs = handle_api_failure{
+              client.jobs.select{|j| j.running? }
+            }
+            break if jobs.empty?
+          end
+          sleep 1
+        end
+      end
+
+      private
+
+      def sql_format(query)
+        tmp = query.dup
+        sql_formatter = AnbtSql::Formatter.new(AnbtSql::Rule.new)
+        sql_formatter.format(tmp).to_s
+      end
+    end
+  end
+end