zillabyte-cli 0.0.9

data/lib/zillabyte/cli/flows.rb

@@ -0,0 +1,702 @@
+require "zillabyte/cli/base"
+require "zillabyte/cli/config"
+require "zillabyte/common"
+require "pty"
+require 'indentation'
+require 'open3'
+require 'securerandom'
+
+# manage custom flows
+#
+class Zillabyte::Command::Flows < Zillabyte::Command::Base
+
+  
+  # flows
+  #
+  #
+  def index
+    self.list
+  end
+
+
+  # flows
+  #
+  # list custom flows
+  #
+  def list
+    
+    headings = ["id", "name", "state"]
+    rows = api.flow.list.map do |row|
+      if headings.size == 0
+        headings = row.keys 
+        headings.delete("rel_dir")
+      end
+      v = row.values_at *headings
+      v
+    end
+    display "flows:\n" + Terminal::Table.new(:headings => headings, :rows => rows).to_s
+    display "Total number of flows: "+rows.length.to_s
+    
+    
+  end
+  
+  alias_command "list", "flows:list"
+
+
+  def status
+    headings = ["name", "implementable", "implemented"]
+    rows = api.flows.list.map do |row|
+      if headings.size == 0
+        headings = row.keys 
+        headings.delete("rel_dir")
+      end
+      v = row.values_at *headings
+      v
+    end
+    display "flow status:\n" + Terminal::Table.new(:headings => headings, :rows => rows).to_s
+    display "Total number of flows in queue: "+rows.length.to_s
+
+  end
+
+  alias_command "status", "flows:status"
+
+  # def status
+  #   headings = ["name", "matches", "emits", "implementable?", "implemented", ]
+  # end
+
+  # flows:push [DIR]
+  #
+  # uploads a flow
+  # 
+  # --config CONFIG_FILE # use the given config file
+  #
+  #Examples:
+  #
+  # $ zillabyte flows:push . 
+  #
+  def push
+    
+    dir = options[:directory] || shift_argument || Dir.pwd
+    res = api.flows.push_directory dir, progress, options
+    
+    if res['error']
+      display "error: #{res['error_message']}"
+    else
+      display "flow ##{res['id']} #{res['action']}"
+    end
+    
+  end
+  alias_command "exec:push", "flows:push"
+  alias_command "push", "flows:push"
+  
+  
+  
+  # flows:pull ID DIR
+  #
+  # pulls a flow source to a directory. The target directory must be empty
+  #
+  # --force # pulls even if the directory exists
+  #
+  #Examples:
+  #
+  # $ zillabyte flows:pull . 
+  #
+  def pull
+    
+    id = options[:id] || shift_argument 
+    dir = options[:directory] || shift_argument
+    
+    error "no id given" if id.nil?
+    error "no directory given" if dir.nil?
+    
+    # Create if not exists..
+    if File.exists?(dir)
+      if Dir.entries(dir).size != 2 and options[:force].nil?
+        error "target directory not empty. use --force to override" 
+      end
+    else
+      FileUtils.mkdir_p(dir)
+    end
+    
+    res = api.flows.pull_to_directory id, dir, progress
+    
+    if res['error']
+      display "error: #{res['error_message']}"
+    else
+      display "flow ##{res['id']} pulled to #{dir}"
+    end
+    
+  end
+  
+  alias_command "pull_to", "flows:pull"
+  alias_command "pull", "flows:pull"
+
+
+
+  
+  # flows:prep [DIR]
+  #
+  # prepares a flow for execution
+  #
+  def prep
+    
+    dir = options[:directory] || shift_argument || Dir.pwd
+    meta = Zillabyte::CLI::Config.get_config_info(dir)
+    
+    case meta["language"]
+    when "ruby"
+      
+      # Execute in the bundler context
+      full_script = File.join(dir, meta["script"])
+      cmd = "cd \"#{meta['home_dir']}\"; unset BUNDLE_GEMFILE; bundle install"
+      exec(cmd)
+      
+    when "python"
+      vDir = "#{meta['home_dir']}/vEnv"
+      lock_file = meta['home_dir']+"/zillabyte_thread_lock_file"
+      if File.exists?(lock_file)
+        sleep(1) while File.exists?(lock_file)
+      else
+        begin
+          cmd = "touch #{lock_file}; virtualenv --clear --system-site-packages #{vDir}; PYTHONPATH=~/zb1/multilang/python/Zillabyte #{vDir}/bin/pip install -r #{meta['home_dir']}/requirements.txt"
+          system cmd, :out => :out
+        ensure
+          File.delete(lock_file)
+        end
+      end
+      
+    when "js"
+    end
+    
+  end
+  alias_command "prep", "flows:prep"
+
+
+
+  
+  # flows:init [LANG] [DIR]
+  #
+  # initializes a new executable in DIR
+  # defaults to a ruby executable for the current directory
+  #
+  #Examples:
+  #
+  # $ zillabyte flows:init python contact_extractor
+  #
+  def init
+    
+    
+    lang = options[:lang] || shift_argument || "ruby"
+    dir = options[:dir] || shift_argument || Dir.pwd
+    languages = ["ruby","python", "js"]
+
+    error "Unsupported language #{lang}. We only support #{languages.join(', ')}." if not languages.include? lang
+
+    display "initializing empty #{lang} flow in #{dir}"
+    FileUtils.cp_r( File.expand_path("../templates/#{lang}", __FILE__) + "/." , dir )
+    
+    
+  end
+  alias_command "exec:init", "flows:init"
+
+  
+  
+
+  # flows:test [TEST_DATASET_ID]
+  #
+  # tests a local flow with sample data
+  # 
+  # --config CONFIG_FILE  # use the given config file
+  # --output OUTPUT_FILE  # writes sink output to a file
+  # --wait MAX             # max time to spend on each operation (default 10 seconds)
+  # --batches BATCHES  # number of batches to emit (default 1)
+  #
+  def test
+    
+    output = options[:output]
+    max_seconds = (options[:wait] || "30").to_i
+    batches = (options[:batches] || "1").to_i
+
+    def read_message(read_stream, color)
+      msg = nil
+      read_stream.each do |line|
+        line.strip!
+        if(line == "end")
+          return msg
+        end
+        begin
+          hash = JSON.parse(line)
+          if(hash["command"] == "done")
+            msg = "done"
+          else
+            msg = hash.to_json
+          end
+        rescue
+          next
+        end
+      end
+      msg
+    end
+
+    def write_message(write_stream, msg)
+      write_stream.write msg.strip + "\n"
+      write_stream.write "end\n"
+      write_stream.flush
+    end
+
+    def handshake(write_stream, read_stream, node, color)
+      begin
+        write_message write_stream, "{\"pidDir\": \"/tmp\"}\n"
+        read_message read_stream, color # Read to "end\n"
+      rescue Exception => e
+        puts "Error handshaking node: #{node}"
+        raise e
+      end
+    end
+
+
+    # INIT
+    test_data = options[:test_data] || shift_argument
+    dir = options[:dir] || Dir.pwd
+
+    meta = Zillabyte::API::Flows.get_rich_meta_info_from_script(dir, self, {:test => true})
+    if meta.nil?
+      error "this is not a valid zillabyte flow directory"
+      exit
+    end    
+
+    # Show the user what we know about their flow... 
+    display "inferring your flow details..."
+    colors = {}
+    describe_flow(meta, colors)
+
+    
+    # Extract the flow's information.. 
+    nodes = meta["nodes"]
+    write_to_next_each = []
+    write_queue = []
+    stream_messages = {}
+    default_stream = "_default"
+    
+    split_branches = false
+    
+    # Iterate all nodes sequentially and invoke them in separate processes... 
+    nodes.each do |node|
+
+      # Init 
+      type = node["type"]
+      name = node["name"]
+      color = colors[name] || :default
+      
+      op_display = lambda do |msg, override_color = nil|
+        display "#{name} - #{msg}".colorize(override_color || color)
+      end
+
+      # A Spout? 
+      if type == "spout"
+
+        # A spout from relation? 
+        if node['matches'] or node["relation"]
+          matches = node['matches'] || (node["relation"]["query"])
+          op_display.call "Grabbing remote data"
+          res = api.query.agnostic(matches)["rows"]
+          if(res.nil? or res.length == 0)
+            raise NameError, "Could not find data that matches your 'matches' clause"
+          end
+          res.each do |tuple|
+            values = {}
+            meta = {}
+            tuple.each do |k, v|
+              if(k == "id")
+                next
+              elsif(k == "confidence" or k == "since" or k == "source")
+                meta[k] = v
+              else
+                values[k] = v
+              end
+            end
+            read_msg = {"tuple" => values, "meta" => meta}.to_json
+            op_display.call "emit tuple: #{values} #{meta}"
+            stream_messages[default_stream] ||= []
+            stream_messages[default_stream] << read_msg
+          end
+          
+          # Done processing... 
+          next 
+          
+        else
+          
+          # A regular spout.. 
+          stream_messages[default_stream] ||= []
+          batches.times do |i|
+            stream_messages[default_stream] << "{\"command\": \"next\"}\n"
+          end
+          
+        end
+        
+      end
+      
+      
+      
+      # A Sink? 
+      if type == "sink"
+       
+        if split_branches || stream_messages.size > 1 
+          sink_stream = node["consumes"]
+          messages = stream_messages[sink_stream] || []
+        else
+          messages = stream_messages.values.first || []
+        end
+        
+        table = Terminal::Table.new :title => name
+        csv_str = CSV.generate do |csv|
+          header_written = false;
+          messages.each do |msg|
+            obj = JSON.parse(msg)
+            if obj['tuple']
+              if header_written == false
+                keys = [obj['tuple'].keys, obj['meta'].keys].flatten
+                csv << keys
+                table << keys
+                table << :separator
+                header_written = true
+              end
+              vals = [obj['tuple'].values, obj['meta'].values].flatten
+              csv << vals
+              table << vals
+            end
+          end
+        end
+        
+        display table.to_s.colorize(color)
+        
+        if output
+          filename = "#{output}.csv"
+          f = File.open(filename, "w")
+          f.write(csv_str)
+          f.close()
+          op_display.call "output written to #{filename}"
+        end
+        
+        
+        
+        next        
+      end
+      
+
+      cmd = command("--execute_live --name #{name}")
+      begin
+        
+        # Start the operation... 
+        op_display.call "beginning #{type} #{name}"
+        Open3.popen3(cmd) do |stdin, stdout, stderr, wait_thread|
+          begin
+            
+            # Init 
+            handshake stdin, stdout, node, color
+            write_queue = []
+            read_queue = []
+          
+            # Get the incoming stream 
+            if !split_branches && stream_messages.size == 1  # i.e. the number of streams we're dealing with right now
+              # Assume default stream 
+              stream_name = stream_messages.keys.first
+              write_queue = stream_messages.values.first.clone
+              stream_messages.delete(stream_name)
+            else
+              # Multiple streams...
+              split_branches = true;
+              if node['consumes'].nil?
+                error "The node #{name} must declare which stream it 'consumes'"
+              end
+              stream_name = node["consumes"]
+              write_queue = stream_messages[stream_name].clone()
+              stream_messages.delete(stream_name)
+            end
+          
+            # Start writing the messages... 
+            stuff_to_read = false
+            writing_thread = Thread.start do 
+              until(write_queue.empty?)
+                
+                # Make sure we're not reading anything... 
+                while(stuff_to_read)   # TODO: semaphores
+                  sleep 0.5 # spin wait
+                end
+                
+                # Get next mesage
+                write_msg = write_queue.shift
+                
+                # Make it human-understable
+                write_json = JSON.parse(write_msg)
+                if write_json['tuple']
+                  op_display.call "receiving: #{write_json['tuple']}"
+                elsif write_json['command'] == 'next'
+                  op_display.call "starting next spout batch"
+                else
+                  puts write_json
+                end
+                
+                # Actually send it to the process 
+                begin
+                  write_message stdin, write_msg
+                  stuff_to_read = true
+                  sleep 0.1
+                rescue Exception => e
+                  puts "Error running #{cmd}: #{e}"
+                  raise e
+                end
+              end
+            end
+            
+            # Start reading messages... 
+            reading_thread = Thread.start do 
+              while(true)
+
+                # Get next message
+                read_msg = read_message(stdout, color)
+                if read_msg == "done" || read_msg.nil?
+                  stuff_to_read = false
+                  if write_queue.empty?
+                    break  # exit while loop
+                  else
+                    sleep 0.5 # spin wait
+                    next
+                  end
+                end
+                stuff_to_read = true
+                
+                # Process message
+                obj = JSON.parse(read_msg)
+                if obj['tuple']
+                  
+                  # Conver to a incoming tuple for the next operation
+                  next_msg = {
+                    :tuple => obj['tuple'],
+                    :meta => obj['meta']
+                  }
+                  emit_stream = obj['stream'] || default_stream
+                  stream_messages[emit_stream] ||= []
+                  stream_messages[emit_stream] << next_msg.to_json
+                  op_display.call "emitted: #{obj['tuple']} to #{emit_stream}"
+                elsif obj['command'] == 'log'
+                  op_display.call "log: #{obj['msg']}"
+                else
+                  error "unknown message: #{read_msg}"
+                end
+                
+              end
+            end
+            
+            
+            # stderr thread 
+            stderr_thread = Thread.start do 
+              stderr.each do |line|
+                op_display.call("stderr: #{line}", :red)
+              end
+            end
+            
+            begin
+              killed = Timeout.timeout(max_seconds) do
+                reading_thread.join()
+                writing_thread.join()
+                stderr_thread.kill()
+                op_display.call "completed #{type} #{name}"
+              end
+            rescue Timeout::Error
+              op_display.call "max time reached. preempting #{type} #{name}. set --wait to increase", :red
+              reading_thread.kill() if reading_thread.alive?
+              writing_thread.kill() if writing_thread.alive?
+              stderr_thread.kill() if stderr_thread.alive?
+            end
+
+          rescue Errno::EIO
+            puts "Errno:EIO error, but this probably just means " +
+                 "that the process has finished giving output"
+          end
+        end
+      rescue PTY::ChildExited
+        puts "The child process exited!"
+      end
+    end
+    
+  end  
+  alias_command "test", "flows:test"
+  
+  
+  
+  
+  
+  # flows:status ID
+  #
+  # gets the flows status
+  # 
+  # --config CONFIG_FILE # use the given config file
+  #
+  def status
+
+    id = options[:id] || shift_argument
+    
+    display api.flows.get(id)
+   
+  end  
+  
+  
+  
+  
+  
+  
+  # flows:kill ID
+  #
+  # kills the given flow
+  # 
+  # --config CONFIG_FILE # use the given config file
+  #
+  def kill
+
+    id = options[:id] || shift_argument
+    api.flows.kill(id)
+    display "flow #{id} killed"
+   
+  end  
+  
+  
+  
+
+
+  
+  # flows:live_run [DIR]
+  #
+  # runs a local flow with live data
+  # 
+  # --config CONFIG_FILE # use the given config file
+  #
+  def live_run
+
+    name = options[:name] || shift_argument
+    thread_id = options[:thread] || shift_argument || ""
+    dir = options[:directory] || shift_argument || Dir.pwd
+    meta = Zillabyte::CLI::Config.get_config_info(dir)
+  
+    if meta.nil?
+      throw "could not find meta information for: #{dir}"
+    end
+   
+    if(thread_id == "")
+      exec(command("--execute_live --name #{name.to_s}"))
+    else 
+      exec(command("--execute_live --name #{name.to_s} --pipe #{thread_id}"))
+    end
+  end  
+  alias_command "live_run", "flows:live_run"
+  
+  
+  # flows:info [DIR]
+  # 
+  # outputs the info for the flow in the dir.
+  # 
+  # --pretty # Pretty prints the info output
+  #
+  def info
+    info_file = SecureRandom.uuid
+    cmd = command("--info --file #{info_file}")
+    flow_info = Zillabyte::Command::Flows.get_info(cmd, info_file)
+
+    if options[:pretty]
+      puts JSON.pretty_generate(JSON.parse(flow_info))
+    else
+      puts flow_info
+    end
+    exit
+  end
+  alias_command "info", "flows:info"
+
+  def self.get_info(cmd, info_file, options = {})
+    response = `#{cmd}`
+    if($?.exitstatus == 1)
+      File.delete("#{info_file}") if File.exists?(info_file)
+      puts "error: #{response}" if options[:test]
+      Process.exit 1
+    end
+
+    flow_info = {}
+    File.open("#{info_file}", "r+").each do |line|
+      line = JSON.parse(line)
+      if(line["type"])
+        flow_info["nodes"] << line
+      else
+        flow_info = line
+        flow_info["nodes"] = []
+      end
+    end
+    File.delete("#{info_file}")
+
+    flow_info = flow_info.to_json
+    if(File.exists?("info_to_java.in"))
+      java_pipe = open("info_to_java.in","w+")
+      java_pipe.puts(flow_info+"\n")
+      java_pipe.flush
+      java_pipe.close()
+    end
+
+    flow_info
+  end
+  
+  private 
+  
+  
+  def command(arg="--execute_live", ignore_stderr = false)
+
+    dir = options[:directory] || shift_argument || Dir.pwd
+    meta = Zillabyte::CLI::Config.get_config_info(dir, progress=nil, options)
+    #meta = Zillabyte::API::Functions.get_rich_meta_info_from_script(dir, self)
+    error "could not extract meta information. missing zillabyte.conf.yml?" if meta.nil? 
+    error meta["error_message"] if meta['status'] == "error"
+    full_script = File.join(dir, meta["script"])
+    stderr_opt = "2> /dev/null" if ignore_stderr
+
+    case meta["language"]
+    when "ruby"
+      # Execute in the bundler context
+      cmd = "cd \"#{dir}\"; unset BUNDLE_GEMFILE; ZILLABYTE_HARNESS=1 bundle exec ruby \"#{full_script}\" #{arg}  #{stderr_opt}"
+      
+    when "python"#{
+      if(File.directory?("#{dir}/vEnv"))
+        cmd = "cd \"#{dir}\"; PYTHONPATH=~/zb1/multilang/python/Zillabyte #{dir}/vEnv/bin/python \"#{full_script}\" #{arg}  #{stderr_opt}"
+      else
+        cmd = "cd \"#{dir}\"; PYTHONPATH=~/zb1/multilang/python/Zillabyte python \"#{full_script}\" #{arg}  #{stderr_opt}"
+      end
+ 
+    when "js"
+#      cmd = "#{Zillabyte::API::CASPERJS_BIN} #{Zillabyte::API::API_CLIENT_JS}  \"#{full_script}\" #{arg}"
+      cmd = "cd \"#{dir}\"; NODE_PATH=~/zb1/multilang/js #{Zillabyte::API::NODEJS_BIN}  \"#{full_script}\" #{arg} #{stderr_opt}"
+
+    else
+      error "no language specified"
+    end
+    
+    return cmd
+
+  end
+  
+  
+  def describe_flow(meta, colors = {})
+    @colors ||= [:green, :yellow, :magenta, :cyan, :light_black, :light_green, :light_yellow, :light_blue, :light_magenta, :light_cyan]
+    rjust = 20
+    display "#{'flow name'.rjust(rjust)}: #{meta['name']}"
+    display "#{'flow language'.rjust(rjust)}: #{meta['language']}"
+    
+    meta['nodes'].each_with_index do |node, index|
+      colors[node['name']] ||= @colors.shift
+      color = colors[node['name']]
+      display (("="*rjust + "  operation ##{index}").colorize(color))
+      display "#{"name".rjust(rjust)}: #{node['name'].to_s.colorize(color)}"
+      display "#{"type".rjust(rjust)}: #{node['type'].to_s.colorize(color)}"
+      display "#{"matches".rjust(rjust)}: #{JSON.pretty_generate(node['matches']).indent(rjust+2).lstrip.colorize(color)}" if node['matches']
+      display "#{"emits".rjust(rjust)}: #{JSON.pretty_generate(node['emits']).indent(rjust+2).lstrip.colorize(color)}" if node['emits']
+    end
+    
+  end
+
+
+end