zillabyte-cli 0.0.24 → 0.1.0

data/lib/zillabyte/runner/component_operation.rb

@@ -0,0 +1,529 @@
+require 'json'
+require 'mkfifo'
+require "zillabyte/runner"
+require "zillabyte/runner/multilang_operation"
+require "zillabyte/runner/component_runner"
+require 'zillabyte/api/components'
+
+
+
+# Emulate Component Operations
+class Zillabyte::Runner::ComponentOperation
+
+  NEXT_MESSAGE = "{\"command\": \"next\"}\n"
+  END_CYCLE_MESSAGE = "{\"command\": \"end_cycle\"}\n"
+  ENDMARKER = "\nend\n"
+
+  # Run the operation
+  def self.run(node, dir, consumee, consumer_pipes, tester, meta, options = {})
+
+
+    @__node = node
+    @__name = node["name"]
+    @__type = node["type"]
+    @__dir = dir
+    @__consumee = consumee
+    @__consumer_pipes = consumer_pipes
+    @__tester = tester
+
+    @__meta = meta
+    @__options = options
+    @__output_type = options[:output_type]
+
+    # Each consumer of a stream gets its own queue and message passing
+    @__emit_queues = {}
+    @__consumer_pipes.each_pair do |stream, consumers|
+      consumers.each_key do |consumer|
+        @__emit_queues[stream] ||= {}
+        @__emit_queues[stream][consumer] = {:write_queue => [], :ready => true}
+      end
+    end
+
+    begin
+      case @__type
+      when "source"
+        self.run_input()
+      when "sink"
+        self.run_output()
+      when "component"
+        self.run_rpc_component()
+      else
+        Zillabyte::Runner::MultilangOperation.run(node, dir, consumee, consumer_pipes, tester, meta, options)
+      end
+    rescue => e
+      cdisplay e.message
+      cdisplay e.backtrace
+    end
+
+  end
+
+
+  # Run a component input
+  def self.run_input()
+    input = @__options[:input]
+
+    # Read input from file
+    if input
+      cdisplay "reading from file...."
+      csv_rows = CSV.read("#{input}")
+
+    else
+
+      stdin = @__consumee[:rd_child]
+      loop do 
+
+        msg = stdin.gets
+
+        # Build tuple
+        begin
+          tuple = JSON.parse(msg)
+        rescue JSON::ParserError 
+          cdisplay "Error: invalid JSON"           
+          next
+        end
+
+        # Check input for correct fields
+        has_fields = true
+        fields = @__node['fields'].flat_map {|h| h.keys}
+        tuple.each_key do |field|
+          if !fields.include? field
+            cdisplay "Error: invalid schema for node"
+            next
+          end
+        end
+        tuple_json = build_tuple_json(tuple)
+        display_json = Hash[JSON.parse(tuple_json)["tuple"].map {|k,v| [truncate_message(k), truncate_message(v)]}].to_json
+        send_to_consumers(tuple_json)   
+      end
+
+    end
+   
+  end
+
+
+  # Send to and manage an RPC component 
+  def self.run_rpc_component()
+
+    # Index streams and consumers by their pipes for lookup
+    consumer_hash = {}
+    @__emit_queues.each_pair do |stream, consumers|
+      consumers.each_key do |consumer|
+        read_stream = @__consumer_pipes[stream][consumer][:rd_parent]
+        consumer_hash[read_stream] = {:stream => stream, :consumer => consumer}
+      end
+    end
+    
+    # Keep track of how many consumers to handle before exiting
+    consumers_running = consumer_hash.keys.length
+
+    # Begin cycle
+    end_cycle_received = false
+
+
+    # TODO multiple inputs
+    # The input component(singular at the moment)
+    read_streams = consumer_hash.keys.concat [@__consumee[:rd_child]]
+
+    # Receive and handle messages
+    loop do
+
+      # Read from a stream
+      rs = select_read_streams(read_streams)
+      rs.each do |r|
+        # Read a message
+        obj = read_message(r)
+
+        # Handle tuple through RPC
+        if obj['tuple']
+
+          display_json = Hash[obj['tuple'].map{|k, v| [truncate_message(k), truncate_message(v)]}].to_json
+          
+
+          # Start communication with API
+          api = "API #{@__tester.session.api}"
+          caller = Zillabyte::API::Components.new(api)
+          component_id = @__node["id"]
+          output_format = @__node["output_format"]
+
+          # Send RPC call
+          tuple = obj['tuple']
+          cdisplay "sending RPC call..."
+          call_options = {
+            :rpc_inputs => [tuple],
+            :output_format => output_format
+          }
+
+          res = caller.rpc(component_id, call_options)
+          cdisplay "received API response : #{res}"
+
+          #TODO handle errors
+          if res['error']
+            cdisplay("error: #{res['error']}")
+            next
+          end
+          run_ids = res['run_ids']
+          status = res['status']
+
+
+          # Send result call
+          loop do
+
+            # If a call for results is not ready, save the execution ID for the next cycle
+            ids_in_progress = []
+            res = caller.get_rpc_results(component_id, {:execute_ids => run_ids})
+            if res['error']
+              cdisplay("error: #{res['error']}")
+              next
+            end
+
+        
+            # check results
+            res['results'].each do |run_id, hash|
+              
+              # We have results
+              if hash['status'] == "complete"
+
+                # Handle Tuple in here
+                data_streams = hash['data']
+
+                # if stream_hash.nil?
+                #   cdisplay("error : no results!")
+                #   next
+                # else
+                #   stream_hash.each_pair
+                # end
+
+                # Handle each resulting tuple
+                data_streams.each_pair do |stream, data_tuples|
+                  data_tuples.each do |data_tuple|
+                    tuple_json = build_tuple_json(data_tuple)
+
+                    #send or enqueue the tuple to all consumers of the stream
+                    @__emit_queues.each_pair do |stream, consumers|
+                      consumers.each_pair do |consumer, emitter|
+                        if emitter[:ready]
+                          emit_consumer_tuple(stream, consumer, tuple_json)
+                        else
+                          @__emit_queues[stream][consumer][:write_queue] << tuple_json
+                        end
+                      end
+                    end
+
+                  end
+                end
+
+              else
+                cdisplay "RPC #{run_id} has not completed... "
+                ids_in_progress << run_id
+              end
+            end
+            run_ids = ids_in_progress
+
+            # If no more IDs to run, we are done
+            if run_ids.empty?
+              break
+            end
+
+            # Dont spam the API
+            sleep(2)
+          end
+
+          # Ask for next tuple
+          write_message(@__consumee[:wr_child], NEXT_MESSAGE)
+
+
+        # End cycle
+        elsif obj['command'] 
+          case obj["command"]
+
+          # Consumer is ready for a message
+          when "next"
+            stream = consumer_hash[r][:stream]
+            consumer = consumer_hash[r][:consumer]
+
+            @__emit_queues[stream][consumer][:ready] = true
+            tuple_json = get_consumer_tuple(stream, consumer)
+
+            # End cycle for consumer if it has processed all tuples
+            if tuple_json.nil? && end_cycle_received
+              write_stream = get_write_stream(stream, consumer)
+              write_message(write_stream, END_CYCLE_MESSAGE)
+              consumers_running -= 1
+              if consumers_running == 0
+                break
+              end
+
+              # TODO break if last consumer
+            elsif !tuple_json.nil?
+              # Emit tuple to consumer
+              emit_consumer_tuple(stream, consumer, tuple_json)
+              emitted = true
+            end
+
+          when "end_cycle"
+            end_cycle_received = true
+          end
+        end
+
+      end
+
+      # Exit after ending consumer cycles
+      if consumers_running == 0
+        break
+      end
+    end
+
+  end
+
+  # Run a component output
+  def self.run_output()
+
+
+    output = @__options[:output]
+    header_written = false
+    read_stream = @__consumee[:rd_child]
+
+    loop do
+      # Read a tuple
+      obj = read_message(read_stream)
+
+
+      # Add row
+      if obj['tuple']
+
+        if output
+          filename = "#{output}.csv"
+          t = obj['tuple']
+          m = obj['meta'] || {}
+
+          # Write header
+          if header_written == false
+            keys = [t.keys, m.keys].flatten
+            csv_str = keys
+            f = File.open(filename, "w")
+            f.write(csv_str)
+            f.close()
+            header_written = true
+          end
+
+          # Write CSV row
+          vals = [t.values, m.values].flatten
+          csv_str = vals.to_csv
+          
+          f = File.open(filename, "w")
+          f.write(csv_str)
+          f.close()
+        else
+
+          display_json = Hash[obj['tuple'].map{|k, v| [truncate_message(k), truncate_message(v)]}].to_json
+          cdisplay "received #{display_json}"
+        end
+
+        # Ready for next message
+        write_message(@__consumee[:wr_child], NEXT_MESSAGE)
+
+      # End cycle
+      elsif obj['command'] && obj['command'] == "end_cycle"
+        break
+      end
+    end
+
+  end
+
+
+  private
+
+  BUFSIZE = 8192
+
+  # Each reading pipe has a read buffer and message queue
+  @__read_buffers = {}
+  @__read_buffered_messages = {}
+
+
+  # Return availible reading streams
+  def self.select_read_streams(read_streams)
+
+    rs = []
+    read_streams.each do |read_stream|
+      @__read_buffered_messages[read_stream] ||= []
+      if !@__read_buffered_messages[read_stream].empty?
+        rs << read_stream
+      end
+    end
+    return rs unless rs.empty?
+
+    rs, ws, es = IO.select(read_streams, [], [])
+    return rs
+  end
+
+
+  # Read a JSON message
+  def self.read_message(read_stream)
+
+    @__read_buffers[read_stream] ||= ""
+    @__read_buffered_messages[read_stream] ||= []
+    if !@__read_buffered_messages[read_stream].empty?
+        obj = @__read_buffered_messages[read_stream].shift
+      return obj
+    end
+
+    # read message from stream
+    loop do
+
+      while !@__read_buffers[read_stream].include? ENDMARKER
+        segment = read_stream.sysread(BUFSIZE)
+        @__read_buffers[read_stream] << segment
+      end
+
+      read_buffer = @__read_buffers[read_stream]
+      if read_buffer.include? ENDMARKER
+        objs = read_buffer.split(ENDMARKER)
+        ends = read_buffer.scan(ENDMARKER)
+        if objs.count == ends.count # We have a full number of messages
+          objs.each do |obj|
+            begin
+              @__read_buffered_messages[read_stream] << JSON.parse(obj)
+            rescue JSON::ParserError 
+              cdisplay "READMESSAGE: invalid JSON #{obj}"
+            end
+          end
+          @__read_buffers[read_stream] = ""
+          return @__read_buffered_messages[read_stream].shift
+        else
+
+          (0..ends.count-1).each do |i|
+            obj = objs[i] 
+            begin
+              @__read_buffered_messages[read_stream] << JSON.parse(obj)
+            rescue JSON::ParserError 
+              cdisplay "READMESSAGE: invalid JSON #{obj}" 
+            end
+          end
+
+          @__read_buffers[read_stream] = objs[ends.count..-1].join(ENDMARKER)
+          return @__read_buffered_messages[read_stream].shift
+        end
+      end
+    end
+  end
+
+
+  # Write JSON message
+  def self.write_message(write_stream, msg)
+    write_msg = msg.strip + ENDMARKER
+    write_stream.write write_msg
+    write_stream.flush
+  end
+
+  # Format a message for display
+  def self.truncate_message(msg)
+    return msg if(!msg.instance_of?(String))
+    t_length = 50 # truncates entries to this length
+    m_length = msg.length
+    msg_out = m_length > t_length ? msg[0..t_length-3]+"..." : msg
+    msg_out
+  end
+
+
+  # Handshake connection to multilang
+  def self.handshake(write_stream, read_stream)
+    begin
+      write_message write_stream, HANDSHAKE_MESSAGE
+      msg = read_message(read_stream)
+    rescue Exception => e
+      cdisplay(e)
+      cdisplay("Error handshaking node")
+      raise e
+    end
+  end
+
+  # Send object to every consumer of the operation, regardless of stream
+  def self.send_to_consumers(json_obj)
+    @__consumer_pipes.each_pair do |stream, consumers|
+      consumers.each_pair do |consumer, pipe| 
+        write_message(pipe[:wr_parent], json_obj)
+        cdisplay "emitted #{json_obj} to #{consumer}"
+      end
+    end
+  end
+
+
+  # Get the write pipe of the stream consumer
+  def self.get_write_stream(stream, consumer)
+     @__consumer_pipes[stream][consumer][:wr_parent]
+  end
+
+
+  # Get tuple for sending to consumer of stream
+  def self.get_consumer_tuple(stream, consumer)
+    @__emit_queues[stream][consumer][:write_queue].shift
+  end
+
+
+  # Emit tuple_json to the consumer of a stream
+  def self.emit_consumer_tuple(stream, consumer, tuple_json)
+    begin
+      display_json = Hash[JSON.parse(tuple_json)["tuple"].map {|k,v| [truncate_message(k), truncate_message(v)]}].to_json
+    rescue JSON::ParserError
+      cdisplay "Error: invalid JSON"
+    end
+    write_stream = get_write_stream(stream, consumer)
+    write_message(write_stream, tuple_json)
+    @__emit_queues[stream][consumer][:ready] = false
+    cdisplay "emitted tuple #{display_json} to #{consumer} "
+  end
+
+
+  # Build a tuple and format into JSON
+  def self.build_tuple_json(tuple, meta = nil, column_aliases = nil)
+    meta ||= {}
+    column_aliases ||= {}
+    values = {}
+    tuple.each do |k, v|
+     if(k == "id")
+       nextx
+     elsif(k == "confidence" or k == "since" or k == "source")
+       meta[k] = v
+     else  
+      values[k] = v
+     end
+    end
+    tuple_json = {"tuple" => values, "meta" => meta, "column_aliases" => column_aliases}.to_json
+    return tuple_json
+  end
+
+
+  # Construct a multilang command
+  def self.command(arg, ignore_stderr=false)
+    cdisplay("could not extract meta information. missing zillabyte.conf.yml?") if @__meta.nil?
+
+    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 = "cd \"#{@__dir}\"; NODE_PATH=~/zb1/multilang/js/src/lib #{Zillabyte::API::NODEJS_BIN}  \"#{full_script}\" #{arg} #{stderr_opt}"
+    else
+      cdisplay("no language specified")
+    end
+    return cmd
+  end
+
+
+  # Display a colored, formatted message
+  def self.cdisplay(msg)
+
+    @__tester.cdisplay(@__name, msg)
+  end
+
+
+end