presto-client 0.5.14 → 0.6.4

data/lib/presto/client/models.rb

@@ -29,7 +29,9 @@ module Presto::Client
   require 'presto/client/model_versions/0.173.rb'
   require 'presto/client/model_versions/0.178.rb'
   require 'presto/client/model_versions/0.205.rb'
+  require 'presto/client/model_versions/303.rb'
+  require 'presto/client/model_versions/316.rb'
 
-  Models = ModelVersions::V0_205
+  Models = ModelVersions::V316
 
 end

data/lib/presto/client/query.rb

@@ -16,6 +16,7 @@
 module Presto::Client
 
   require 'faraday'
+  require 'faraday_middleware'
   require 'presto/client/models'
   require 'presto/client/errors'
   require 'presto/client/faraday_client'
@@ -125,16 +126,13 @@ module Presto::Client
     end
 
     def close
-      @api.cancel_leaf_stage
+      @api.close
       nil
     end
 
     def raise_if_failed
-      if @api.closed?
+      if @api.client_aborted?
         raise PrestoClientError, "Query aborted by user"
-      elsif @api.exception?
-        # query is gone
-        raise @api.exception
       elsif @api.query_failed?
         results = @api.current_results
         error = results.error

data/lib/presto/client/statement_client.rb

@@ -31,8 +31,7 @@ module Presto::Client
 
       @options = options
       @query = query
-      @closed = false
-      @exception = nil
+      @state = :running
       @retry_timeout = options[:retry_timeout] || 120
       if model_version = @options[:model_version]
         @models = ModelVersions.const_get("V#{model_version.gsub(".", "_")}")
@@ -76,7 +75,7 @@ module Presto::Client
 
       # TODO error handling
       if response.status != 200
-        raise PrestoHttpError.new(response.status, "Failed to start query: #{response.body} (#{response.status})")
+        exception! PrestoHttpError.new(response.status, "Failed to start query: #{response.body} (#{response.status})")
       end
 
       @results_headers = response.headers
@@ -91,14 +90,20 @@ module Presto::Client
       !!@options[:debug]
     end
 
-    def closed?
-      @closed
+    def running?
+      @state == :running
     end
 
-    attr_reader :exception
+    def client_aborted?
+      @state == :client_aborted
+    end
+
+    def client_error?
+      @state == :client_error
+    end
 
-    def exception?
-      @exception
+    def finished?
+      @state == :finished
     end
 
     def query_failed?
@@ -106,7 +111,7 @@ module Presto::Client
     end
 
     def query_succeeded?
-      @results.error == nil && !@exception && !@closed
+      @results.error == nil && finished?
     end
 
     def current_results
@@ -117,16 +122,29 @@ module Presto::Client
       @results_headers
     end
 
+    def query_id
+      @results.id
+    end
+
     def has_next?
       !!@results.next_uri
     end
 
+    def exception!(e)
+      @state = :client_error
+      raise e
+    end
+
     def advance
-      if closed? || !has_next?
+      return false unless running?
+
+      unless has_next?
+        @state = :finished
         return false
       end
 
       uri = @results.next_uri
+
       response = faraday_get_with_retry(uri)
       @results_headers = response.headers
       @results = decode_model(uri, parse_body(response), @models::QueryResults)
@@ -150,8 +168,7 @@ module Presto::Client
         if body.size > 1024 + 3
           body = "#{body[0, 1024]}..."
         end
-        @exception = PrestoHttpError.new(500, "Presto API returned unexpected structure at #{uri}. Expected #{body_class} but got #{body}: #{e}")
-        raise @exception
+        exception! PrestoHttpError.new(500, "Presto API returned unexpected structure at #{uri}. Expected #{body_class} but got #{body}: #{e}")
       end
     end
 
@@ -166,8 +183,7 @@ module Presto::Client
           JSON.parse(response.body, opts = JSON_OPTIONS)
         end
       rescue => e
-        @exception = PrestoHttpError.new(500, "Presto API returned unexpected data format. #{e}")
-        raise @exception
+        exception! PrestoHttpError.new(500, "Presto API returned unexpected data format. #{e}")
       end
     end
 
@@ -184,8 +200,7 @@ module Presto::Client
           # temporally error to retry
           response = nil
         rescue => e
-          @exception = e
-          raise @exception
+          exception! e
         end
 
         if response
@@ -195,8 +210,7 @@ module Presto::Client
 
           if response.status != 503  # retry only if 503 Service Unavailable
             # deterministic error
-            @exception = PrestoHttpError.new(response.status, "Presto API error at #{uri} returned #{response.status}: #{response.body}")
-            raise @exception
+            exception! PrestoHttpError.new(response.status, "Presto API error at #{uri} returned #{response.status}: #{response.body}")
           end
         end
 
@@ -204,18 +218,14 @@ module Presto::Client
 
         attempts += 1
         sleep attempts * 0.1
-      end while (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) < @retry_timeout && !@closed
+      end while (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) < @retry_timeout && !client_aborted?
 
-      @exception = PrestoHttpError.new(408, "Presto API error due to timeout")
-      raise @exception
+      exception! PrestoHttpError.new(408, "Presto API error due to timeout")
     end
 
     def raise_if_timeout!
       if @started_at
-        if @results && @results.next_uri == nil
-          # query is already done
-          return
-        end
+        return if finished?
 
         elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - @started_at
 
@@ -234,30 +244,34 @@ module Presto::Client
 
     def raise_timeout_error!
       if query_id = @results && @results.id
-        raise PrestoQueryTimeoutError, "Query #{query_id} timed out"
+        exception! PrestoQueryTimeoutError.new("Query #{query_id} timed out")
       else
-        raise PrestoQueryTimeoutError, "Query timed out"
+        exception! PrestoQueryTimeoutError.new("Query timed out")
       end
     end
 
     def cancel_leaf_stage
-      if uri = @results.next_uri
-        response = @faraday.delete do |req|
+      if uri = @results.partial_cancel_uri
+        @faraday.delete do |req|
           req.url uri
         end
-        return response.status / 100 == 2
       end
-      return false
     end
 
     def close
-      return if @closed
+      return unless running?
 
-      # cancel running statement
-      # TODO make async reqeust and ignore response?
-      cancel_leaf_stage
+      @state = :client_aborted
+
+      begin
+        if uri = @results.next_uri
+          @faraday.delete do |req|
+            req.url uri
+          end
+        end
+      rescue => e
+      end
 
-      @closed = true
       nil
     end
   end

data/lib/presto/client/version.rb

@@ -15,6 +15,6 @@
 #
 module Presto
   module Client
-    VERSION = "0.5.14"
+    VERSION = "0.6.4"
   end
 end

data/modelgen/model_versions.rb

@@ -131,6 +131,7 @@ module Presto::Client::ModelVersions
           when "apply"              then ApplyNode
           when "assignUniqueId"     then AssignUniqueId
           when "lateralJoin"        then LateralJoinNode
+          when "statisticsWriterNode" then StatisticsWriterNode
         end
         if model_class
            node = model_class.decode(hash)
@@ -197,9 +198,25 @@ module Presto::Client::ModelVersions
         end
         obj = allocate
         model_class = case hash["@type"]
-            when "CreateHandle"       then CreateHandle
-            when "InsertHandle"       then InsertHandle
-            when "DeleteHandle"       then DeleteHandle
+            when "CreateTarget"       then CreateTarget
+            when "InsertTarget"       then InsertTarget
+            when "DeleteTarget"       then DeleteTarget
+        end
+        if model_class
+           model_class.decode(hash)
+        end
+      end
+    end
+
+    class << WriteStatisticsTarget =
+        Base.new(:type, :handle)
+      def decode(hash)
+        unless hash.is_a?(Hash)
+          raise TypeError, "Can't convert #{hash.class} to Hash"
+        end
+        obj = allocate
+        model_class = case hash["@type"]
+            when "WriteStatisticsHandle"       then WriteStatisticsHandle
         end
         if model_class
            model_class.decode(hash)
@@ -246,6 +263,13 @@ module Presto::Client::ModelVersions
       end
     end
 
+    class ResourceGroupId < Array
+      def initialize(array)
+        super()
+        concat(array)
+      end
+    end
+
     ##
     # Those model classes are automatically generated
     #

data/modelgen/modelgen.rb

@@ -13,13 +13,13 @@ erb = ERB.new(File.read(template_path))
 
 source_path = source_dir
 
-predefined_simple_classes = %w[StageId TaskId Lifespan ConnectorSession]
-predefined_models = %w[DistributionSnapshot PlanNode EquiJoinClause WriterTarget OperatorInfo HashCollisionsInfo]
+predefined_simple_classes = %w[StageId TaskId Lifespan ConnectorSession ResourceGroupId]
+predefined_models = %w[DistributionSnapshot PlanNode EquiJoinClause WriterTarget WriteStatisticsTarget OperatorInfo HashCollisionsInfo]
 
-assume_primitive = %w[Object Type Long Symbol QueryId PlanNodeId PlanFragmentId MemoryPoolId TransactionId URI Duration DataSize DateTime ColumnHandle ConnectorTableHandle ConnectorOutputTableHandle ConnectorIndexHandle ConnectorColumnHandle ConnectorInsertTableHandle ConnectorTableLayoutHandle Expression FunctionCall TimeZoneKey Locale TypeSignature Frame TupleDomain<ColumnHandle> SerializableNativeValue ConnectorTransactionHandle OutputBufferId ConnectorPartitioningHandle NullableValue ConnectorId HostAddress JsonNode Node]
-enum_types = %w[QueryState StageState TaskState QueueState PlanDistribution OutputPartitioning Step SortOrder BufferState NullPartitioning BlockedReason ParameterKind FunctionKind PartitionFunctionHandle Scope ErrorType DistributionType PipelineExecutionStrategy JoinType]
+assume_primitive = %w[Object Type Long Symbol QueryId PlanNodeId PlanFragmentId MemoryPoolId TransactionId URI Duration DataSize DateTime ColumnHandle ConnectorTableHandle ConnectorOutputTableHandle ConnectorIndexHandle ConnectorColumnHandle ConnectorInsertTableHandle ConnectorTableLayoutHandle Expression FunctionCall TimeZoneKey Locale TypeSignature Frame TupleDomain<ColumnHandle> SerializableNativeValue ConnectorTransactionHandle OutputBufferId ConnectorPartitioningHandle NullableValue ConnectorId HostAddress JsonNode Node CatalogName QualifiedObjectName]
+enum_types = %w[QueryState StageState TaskState QueueState PlanDistribution OutputPartitioning Step SortOrder BufferState NullPartitioning BlockedReason ParameterKind FunctionKind PartitionFunctionHandle Scope ErrorType DistributionType PipelineExecutionStrategy JoinType ExchangeNode.Type ColumnStatisticType TableStatisticType StageExecutionStrategy SemanticErrorCode]
 
-root_models = %w[QueryResults QueryInfo] + %w[
+root_models = %w[QueryResults QueryInfo BasicQueryInfo] + %w[
 OutputNode
 ProjectNode
 TableScanNode
@@ -42,7 +42,6 @@ IndexJoinNode
 IndexSourceNode
 TableWriterNode
 DeleteNode
-MetadataDeleteNode
 TableFinishNode
 UnnestNode
 ExchangeNode
@@ -54,6 +53,7 @@ ExplainAnalyzeNode
 ApplyNode
 AssignUniqueId
 LateralJoinNode
+StatisticsWriterNode
 ] + %w[
 ExchangeClientStatus
 LocalExchangeBufferInfo
@@ -62,7 +62,8 @@ SplitOperatorInfo
 PartitionedOutputInfo
 JoinOperatorInfo
 WindowInfo
-TableWriterInfo]
+TableWriterInfo
+]
 
 name_mapping = Hash[*%w[
 StatementStats StageStats ClientStageStats
@@ -71,13 +72,14 @@ QueryResults Column ClientColumn
 ].each_slice(3).map { |x, y, z| [[x,y], z] }.flatten(1)]
 
 path_mapping = Hash[*%w[
-ClientColumn presto-client/src/main/java/com/facebook/presto/client/Column.java
-ClientStageStats presto-client/src/main/java/com/facebook/presto/client/StageStats.java
-Column presto-main/src/main/java/com/facebook/presto/execution/Column.java
-QueryStats presto-main/src/main/java/com/facebook/presto/execution/QueryStats.java
-StageStats presto-main/src/main/java/com/facebook/presto/execution/StageStats.java
-PartitionedOutputInfo presto-main/src/main/java/com/facebook/presto/operator/PartitionedOutputOperator.java
-TableWriterInfo presto-main/src/main/java/com/facebook/presto/operator/TableWriterOperator.java
+ClientColumn presto-client/src/main/java/io/prestosql/client/Column.java
+ClientStageStats presto-client/src/main/java/io/prestosql/client/StageStats.java
+Column presto-main/src/main/java/io/prestosql/execution/Column.java
+QueryStats presto-main/src/main/java/io/prestosql/execution/QueryStats.java
+StageStats presto-main/src/main/java/io/prestosql/execution/StageStats.java
+PartitionedOutputInfo presto-main/src/main/java/io/prestosql/operator/PartitionedOutputOperator.java
+TableWriterInfo presto-main/src/main/java/io/prestosql/operator/TableWriterOperator.java
+TableInfo presto-main/src/main/java/io/prestosql/execution/TableInfo.java
 ].map.with_index { |v,i| i % 2 == 0 ? v : (source_path + "/" + v) }]
 
 # model => [ [key,nullable,type], ... ]

data/modelgen/presto_models.rb

@@ -3,13 +3,13 @@ module PrestoModels
   require 'find'
   require 'stringio'
 
-  PRIMITIVE_TYPES = %w[String boolean long int short byte double float Integer]
+  PRIMITIVE_TYPES = %w[String boolean long int short byte double float Integer Double Boolean]
   ARRAY_PRIMITIVE_TYPES = PRIMITIVE_TYPES.map { |t| "#{t}[]" }
 
   class Model < Struct.new(:name, :fields)
   end
 
-  class Field < Struct.new(:key, :nullable, :array, :map, :type, :base_type, :map_value_base_type)
+  class Field < Struct.new(:key, :nullable, :array, :map, :type, :base_type, :map_value_base_type, :base_type_alias)
     alias_method :nullable?, :nullable
     alias_method :array?, :array
     alias_method :map?, :map
@@ -47,10 +47,12 @@ module PrestoModels
 
     private
 
-    PROPERTY_PATTERN = /@JsonProperty\(\"(\w+)\"\)\s+(@Nullable\s+)?([\w\<\>\[\]\,\s]+)\s+\w+/
+    PROPERTY_PATTERN = /@JsonProperty\(\"(\w+)\"\)\s+(@Nullable\s+)?([\w\<\>\[\]\,\s\.]+)\s+\w+/
     CREATOR_PATTERN = /@JsonCreator[\s]+public[\s]+(static\s+)?(\w+)[\w\s]*\((?:\s*#{PROPERTY_PATTERN}\s*,?)+\)/
+    GENERIC_PATTERN = /(\w+)\<(\w+)\>/
 
-    def analyze_fields(model_name, creator_block)
+    def analyze_fields(model_name, creator_block, generic: nil)
+      model_name = "#{model_name}_#{generic}" if generic
       extra = @extra_fields[model_name] || []
       fields = creator_block.scan(PROPERTY_PATTERN).concat(extra).map do |key,nullable,type|
         map = false
@@ -63,7 +65,7 @@ module PrestoModels
           base_type = m[1]
           map_value_base_type = m[2]
           map = true
-        elsif m = /Optional<([\w\[\]]+)>/.match(type)
+        elsif m = /Optional<([\w\[\]\<\>]+)>/.match(type)
           base_type = m[1]
           nullable = true
         elsif m = /OptionalInt/.match(type)
@@ -72,6 +74,9 @@ module PrestoModels
         elsif m = /OptionalLong/.match(type)
           base_type = 'Long'
           nullable = true
+        elsif m = /OptionalDouble/.match(type)
+          base_type = 'Double'
+          nullable = true
         elsif type =~ /\w+/
           base_type = type
         else
@@ -79,7 +84,16 @@ module PrestoModels
         end
         base_type = @name_mapping[[model_name, base_type]] || base_type
         map_value_base_type = @name_mapping[[model_name, map_value_base_type]] || map_value_base_type
-        Field.new(key, !!nullable, array, map, type, base_type, map_value_base_type)
+        
+        if generic
+          base_type = generic if base_type == 'T'
+          map_value_base_type = generic if map_value_base_type == 'T'  
+        end
+        if m = GENERIC_PATTERN.match(base_type)
+          base_type_alias = "#{m[1]}_#{m[2]}"
+        end
+
+        Field.new(key, !!nullable, array, map, type, base_type, map_value_base_type, base_type_alias)
       end
 
       @models[model_name] = Model.new(model_name, fields)
@@ -92,9 +106,15 @@ module PrestoModels
       return fields
     end
 
-    def analyze_model(model_name, parent_model = nil)
+    def analyze_model(model_name, parent_model=  nil, generic: nil)
       return if @models[model_name] || @ignore_types.include?(model_name)
 
+      if m = GENERIC_PATTERN.match(model_name)
+        analyze_model(m[1], generic: m[2])
+        analyze_model(m[2])
+        return
+      end
+
       path = find_class_file(model_name, parent_model)
       java = File.read(path)
 
@@ -114,7 +134,7 @@ module PrestoModels
         fields = analyze_fields(inner_model_name, m[0])
       end
 
-      fields = analyze_fields(model_name, body)
+      fields = analyze_fields(model_name, body, generic: generic)
 
     rescue => e
       puts "Skipping model #{parent_model}/#{model_name}: #{e}"
@@ -214,7 +234,7 @@ module PrestoModels
             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}\"]")
+            expr << convert_expression(field.type, field.base_type_alias || field.base_type, "hash[\"#{field.key}\"]")
           end
         end