presto-client 0.3.3 → 0.4.0

data/lib/presto/client/query.rb

@@ -36,14 +36,20 @@ module Presto::Client
       new StatementClient.new(faraday, query, options)
     end
 
-    def initialize(client)
-      @client = client
+    def initialize(api)
+      @api = api
     end
 
-    attr_reader :client
+    def current_results
+      @api.current_results
+    end
+
+    def advance
+      @api.advance
+    end
 
     def wait_for_data
-      while @client.current_results.data == nil && @client.advance
+      while @api.current_results.data == nil && @api.advance
       end
     end
 
@@ -52,9 +58,9 @@ module Presto::Client
     def columns
       wait_for_data
 
-      raise_error unless @client.query_succeeded?
+      raise_error unless @api.query_succeeded?
 
-      return @client.current_results.columns
+      return @api.current_results.columns
     end
 
     def rows
@@ -74,36 +80,40 @@ module Presto::Client
     def each_row_chunk(&block)
       wait_for_data
 
-      raise_error unless @client.query_succeeded?
+      raise_error unless @api.query_succeeded?
 
       if self.columns == nil
-        raise PrestoError, "Query #{@client.current_results.id} has no columns"
+        raise PrestoError, "Query #{@api.current_results.id} has no columns"
       end
 
       begin
-        if data = @client.current_results.data
+        if data = @api.current_results.data
           block.call(data)
         end
-      end while @client.advance
+      end while @api.advance
+    end
+
+    def query_info
+      @api.query_info
     end
 
     def cancel
-      @client.cancel_leaf_stage
+      @api.cancel_leaf_stage
     end
 
     def close
-      @client.cancel_leaf_stage
+      @api.cancel_leaf_stage
       nil
     end
 
     def raise_error
-      if @client.closed?
+      if @api.closed?
         raise PrestoClientError, "Query aborted by user"
-      elsif @client.exception?
+      elsif @api.exception?
         # query is gone
-        raise @client.exception
-      elsif @client.query_failed?
-        results = @client.current_results
+        raise @api.exception
+      elsif @api.query_failed?
+        results = @api.current_results
         error = results.error
         raise PrestoQueryError.new("Query #{results.id} failed: #{error.message}", results.id, error.error_code, error.failure_info)
       end

data/lib/presto/client/statement_client.rb

@@ -89,7 +89,7 @@ module Presto::Client
 
       body = response.body
       hash = MultiJson.load(body)
-      @results = QueryResults.decode_hash(hash)
+      @results = Models::QueryResults.decode(hash)
     end
 
     private :post_query_request!
@@ -132,27 +132,36 @@ module Presto::Client
       end
       uri = @results.next_uri
 
+      body = faraday_get_with_retry(uri)
+      @results = Models::QueryResults.decode(MultiJson.load(body))
+
+      return true
+    end
+
+    def query_info
+      body = faraday_get_with_retry("/v1/query/#{@results.id}")
+      Models::QueryInfo.decode(MultiJson.load(body))
+    end
+
+    def faraday_get_with_retry(uri, &block)
       start = Time.now
       attempts = 0
 
       begin
         begin
-          response = @faraday.get do |req|
-            req.url uri
-          end
+          response = @faraday.get(uri)
         rescue => e
           @exception = e
           raise @exception
         end
 
         if response.status == 200 && !response.body.to_s.empty?
-          @results = QueryResults.decode_hash(MultiJson.load(response.body))
-          return true
+          return response.body
         end
 
-        if response.status != 503  # retry on 503 Service Unavailable
+        if response.status != 503  # retry only if 503 Service Unavailable
           # deterministic error
-          @exception = PrestoHttpError.new(response.status, "Error fetching next at #{uri} returned #{response.status}: #{response.body}")
+          @exception = PrestoHttpError.new(response.status, "Presto API error at #{uri} returned #{response.status}: #{response.body}")
           raise @exception
         end
 
@@ -160,7 +169,7 @@ module Presto::Client
         sleep attempts * 0.1
       end while (Time.now - start) < 2*60*60 && !@closed
 
-      @exception = PrestoHttpError.new(408, "Error fetching next due to timeout")
+      @exception = PrestoHttpError.new(408, "Presto API error due to timeout")
       raise @exception
     end
 

data/lib/presto/client/version.rb

@@ -15,6 +15,6 @@
 #
 module Presto
   module Client
-    VERSION = "0.3.3"
+    VERSION = "0.4.0"
   end
 end

data/modelgen/modelgen.rb

@@ -0,0 +1,69 @@
+
+if ARGV.length != 3
+  puts "usage: <presto-source-dir> <template.erb> <output.rb>"
+end
+
+source_dir, template_path, output_path = *ARGV
+
+require_relative 'presto_models'
+
+require 'erb'
+erb = ERB.new(File.read(template_path))
+
+source_path = "/Users/frsyuki/project/presto-client-ruby/presto"
+
+predefined_simple_classes = %w[QueryId StageId TaskId PlanNodeId PlanFragmentId ConnectorSession]
+predefined_models = %w[DistributionSnapshot PlanNode]
+
+assume_primitive = %w[Object Type Symbol URI Duration DataSize DateTime ConnectorTableHandle ConnectorOutputTableHandle ConnectorIndexHandle ConnectorColumnHandle Expression FunctionCall]
+enum_types = %w[QueryState StageState TaskState QueueState PlanDistribution OutputPartitioning Step SortOrder]
+
+root_models = %w[QueryResults QueryInfo] + %w[
+OutputNode
+ProjectNode
+TableScanNode
+ValuesNode
+AggregationNode
+MarkDistinctNode
+MaterializeSampleNode
+FilterNode
+WindowNode
+LimitNode
+DistinctLimitNode
+TopNNode
+SampleNode
+SortNode
+ExchangeNode
+SinkNode
+JoinNode
+SemiJoinNode
+IndexJoinNode
+IndexSourceNode
+TableWriterNode
+TableCommitNode
+]
+
+analyzer = PrestoModels::ModelAnalyzer.new(
+  source_path,
+  skip_models: predefined_models + predefined_simple_classes + assume_primitive + enum_types
+)
+analyzer.analyze(root_models)
+models = analyzer.models
+skipped_models = analyzer.skipped_models
+
+formatter = PrestoModels::ModelFormatter.new(
+  base_indent_count: 2,
+  struct_class: "Base",
+  special_struct_initialize_method: "initialize_struct",
+  primitive_types: assume_primitive,
+  skip_types: skipped_models,
+  simple_classes: predefined_simple_classes,
+  enum_types: enum_types,
+)
+formatter.format(models)
+
+@contents = formatter.contents
+
+data = erb.result
+File.open(output_path, 'w') {|f| f.write data }
+

data/modelgen/models.rb

@@ -0,0 +1,132 @@
+#
+# Presto client for Ruby
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+module Presto::Client
+
+  ####
+  ## lib/presto/client/models.rb is automatically generated using "rake modelgen" command.
+  ## You should not edit this file directly. To modify the class definitions, edit
+  ## modelgen/models.rb file and run "rake modelgen".
+  ##
+
+  module Models
+    class Base < Struct
+      class << self
+        alias_method :new_struct, :new
+
+        def new(*args)
+          new_struct(*args) do
+            # make it immutable
+            undef_method :"[]="
+            members.each do |m|
+              undef_method :"#{m}="
+            end
+
+            # replace constructor to receive hash instead of array
+            alias_method :initialize_struct, :initialize
+
+            def initialize(params={})
+              initialize_struct(*members.map {|m| params[m] })
+            end
+          end
+        end
+      end
+    end
+
+    class QueryId < String
+    end
+
+    class StageId < String
+    end
+
+    class TaskId < String
+    end
+
+    class PlanNodeId < String
+    end
+
+    class PlanFragmentId < String
+    end
+
+    class ConnectorSession < Hash
+      def initialize(hash)
+        super()
+        merge!(hash)
+      end
+    end
+
+    module PlanNode
+      def self.decode(hash)
+        model_class = case hash["type"]
+          when "output"             then OutputNode
+          when "project"            then ProjectNode
+          when "tablescan"          then TableScanNode
+          when "values"             then ValuesNode
+          when "aggregation"        then AggregationNode
+          when "markDistinct"       then MarkDistinctNode
+          when "materializeSample"  then MaterializeSampleNode
+          when "filter"             then FilterNode
+          when "window"             then WindowNode
+          when "limit"              then LimitNode
+          when "distinctlimit"      then DistinctLimitNode
+          when "topn"               then TopNNode
+          when "sample"             then SampleNode
+          when "sort"               then SortNode
+          when "exchange"           then ExchangeNode
+          when "sink"               then SinkNode
+          when "join"               then JoinNode
+          when "semijoin"           then SemiJoinNode
+          when "indexjoin"          then IndexJoinNode
+          when "indexsource"        then IndexSourceNode
+          when "tablewriter"        then TableWriterNode
+          when "tablecommit"        then TableCommitNode
+          end
+        model_class.decode(hash) if model_class
+      end
+    end
+
+    # io.airlift.stats.Distribution.DistributionSnapshot
+    class << DistributionSnapshot =
+        Base.new(:max_error, :count, :total, :p01, :p05, :p10, :p25, :p50, :p75, :p90, :p95, :p99, :min, :max)
+      def decode(hash)
+        obj = allocate
+        obj.send(:initialize_struct,
+          hash["maxError"],
+          hash["count"],
+          hash["total"],
+          hash["p01"],
+          hash["p05"],
+          hash["p10"],
+          hash["p25"],
+          hash["p50"],
+          hash["p75"],
+          hash["p90"],
+          hash["p95"],
+          hash["p99"],
+          hash["min"],
+          hash["max"],
+        )
+        obj
+      end
+    end
+
+
+    ##
+    # Those model classes are automatically generated
+    #
+
+<%= @contents %>
+  end
+end

data/modelgen/presto_models.rb

@@ -0,0 +1,210 @@
+
+module PrestoModels
+  require 'find'
+  require 'stringio'
+
+  PRIMITIVE_TYPES = %w[String boolean long int short byte double float]
+
+  class Model < Struct.new(:name, :fields)
+  end
+
+  class Field < Struct.new(:key, :nullable, :array, :map, :type, :base_type, :map_value_base_type)
+    alias_method :nullable?, :nullable
+    alias_method :array?, :array
+    alias_method :map?, :map
+
+    def name
+      @name ||= key.gsub(/[A-Z]/) {|f| "_#{f.downcase}" }
+    end
+  end
+
+  class ModelAnalysisError < StandardError
+  end
+
+  class ModelAnalyzer
+    def initialize(source_path, options={})
+      @source_path = source_path
+      @ignore_types = PRIMITIVE_TYPES + (options[:skip_models] || [])
+      @models = {}
+      @skipped_models = []
+    end
+
+    attr_reader :skipped_models
+
+    def models
+      @models.values.sort_by {|model| model.name }
+    end
+
+    def analyze(root_models)
+      root_models.each {|model_name|
+        analyze_model(model_name)
+      }
+    end
+
+    private
+
+    PROPERTY_PATTERN = /@JsonProperty\(\"(\w+)\"\)\s+(@Nullable\s+)?([\w\<\>\,\s]+)\s+(\w+)/
+    CREATOR_PATTERN = /@JsonCreator[\w\s]+\((?:\s*#{PROPERTY_PATTERN}\s*,?)+\)/
+
+    def analyze_model(model_name)
+      return if @models[model_name] || @ignore_types.include?(model_name)
+
+      path = find_class_file(model_name)
+      java = File.read(path)
+
+      m = CREATOR_PATTERN.match(java)
+      unless m
+        raise ModelAnalysisError, "Can't find JsonCreator of a model class #{model_name}"
+      end
+
+      fields = m[0].scan(PROPERTY_PATTERN).map do |key,nullable,type,field|
+        map = false
+        array = false
+        nullable = !!nullable
+        if m = /(?:List|Set)<(\w+)>/.match(type)
+          base_type = m[1]
+          array = true
+        elsif m = /(?:Map)<(\w+),\s*(\w+)>/.match(type)
+          base_type = m[1]
+          map_value_base_type = m[2]
+          map = true
+        elsif m = /Optional<(\w+)>/.match(type)
+          base_type = m[1]
+          nullable = true
+        elsif type =~ /\w+/
+          base_type = type
+        else
+          raise ModelAnalysisError, "Unsupported type #{type} in model #{model_name}"
+        end
+        Field.new(key, !!nullable, array, map, type, base_type, map_value_base_type)
+      end
+
+      @models[model_name] = Model.new(model_name, fields)
+
+      # recursive call
+      fields.each do |field|
+        analyze_model(field.base_type)
+        analyze_model(field.map_value_base_type) if field.map_value_base_type
+      end
+
+    rescue => e
+      puts "Skipping model #{model_name}: #{e}"
+      @skipped_models << model_name
+    end
+
+    def find_class_file(model_name)
+      @source_files ||= Find.find(@source_path).to_a
+      pattern = /\/#{model_name}.java$/
+      matched = @source_files.find {|path| path =~ pattern }
+      unless matched
+        raise ModelAnalysisError, "Model class #{model_name} is not found"
+      end
+      return matched
+    end
+  end
+
+  class ModelFormatter
+    def initialize(options={})
+      @indent = options[:indent] || '  '
+      @base_indent_count = options[:base_indent_count] || 0
+      @struct_class = options[:struct_class] || 'Struct'
+      @special_struct_initialize_method = options[:special_struct_initialize_method]
+      @primitive_types = PRIMITIVE_TYPES + (options[:primitive_types] || [])
+      @skip_types = options[:skip_types] || []
+      @simple_classes = options[:simple_classes]
+      @enum_types = options[:enum_types]
+      @special_types = options[:special_types] || {}
+      @data = StringIO.new
+    end
+
+    def contents
+      @data.string
+    end
+
+    def format(models)
+      @models = models
+      models.each do |model|
+        @model = model
+
+        puts_with_indent 0, "class << #{model.name} ="
+        puts_with_indent 2, "#{@struct_class}.new(#{model.fields.map {|f| ":#{f.name}" }.join(', ')})"
+        format_decode
+        puts_with_indent 0, "end"
+        line
+      end
+    end
+
+    private
+
+    def line
+      @data.puts ""
+    end
+
+    def puts_with_indent(n, str)
+      @data.puts "#{@indent * (@base_indent_count + n)}#{str}"
+    end
+
+    def format_decode
+      puts_with_indent 1, "def decode(hash)"
+      if @special_struct_initialize_method
+        puts_with_indent 2, "obj = allocate"
+        puts_with_indent 2, "obj.send(:#{@special_struct_initialize_method},"
+      else
+        puts_with_indent 2, "new("
+      end
+
+      @model.fields.each do |field|
+        next if @skip_types.include?(field.base_type) || @skip_types.include?(field.map_value_base_type)
+
+        if @primitive_types.include?(field.base_type) && !field.map?
+          expr = "hash[\"#{field.key}\"]"
+        else
+          expr = ""
+          expr << "hash[\"#{field.key}\"] && " #if field.nullable?
+
+          if field.map?
+            key_expr = convert_expression(field.base_type, field.base_type, "k")
+            value_expr = convert_expression(field.map_value_base_type, field.map_value_base_type, "v")
+            if key_expr == 'k' && value_expr == 'v'
+              expr = "hash[\"#{field.key}\"]"
+            else
+              expr << "Hash[hash[\"#{field.key}\"].to_a.map! {|k,v| [#{key_expr}, #{value_expr}] }]"
+            end
+          elsif field.array?
+            elem_expr = convert_expression(field.base_type, field.base_type, "h")
+            expr << "hash[\"#{field.key}\"].map {|h| #{elem_expr} }"
+          else
+            expr << convert_expression(field.type, field.base_type, "hash[\"#{field.key}\"]")
+          end
+        end
+
+        #comment = "# #{field.base_type}#{field.array? ? '[]' : ''} #{field.key}"
+        #puts_with_indent 3, "#{expr},  #{comment}"
+        puts_with_indent 3, "#{expr},"
+      end
+
+      puts_with_indent 2, ")"
+
+      if @special_struct_initialize_method
+        puts_with_indent 2, "obj"
+      end
+
+      puts_with_indent 1, "end"
+    end
+
+    def convert_expression(type, base_type, key)
+      if @special_types[type]
+        special.call(key)
+      elsif @enum_types.include?(type)
+        "#{key}.downcase.to_sym"
+      elsif @primitive_types.include?(base_type)
+        key
+      elsif @simple_classes.include?(base_type)
+        "#{base_type}.new(#{key})"
+      else  # model class
+        "#{base_type}.decode(#{key})"
+      end
+    end
+  end
+end
+