presto-client-legacy 0.4.17

data/lib/presto/client/query.rb

@@ -0,0 +1,176 @@
+#
+# 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 PrestoLegacy::Client
+
+  require 'faraday'
+  require 'presto/client/models'
+  require 'presto/client/errors'
+  require 'presto/client/statement_client'
+
+  class Query
+    def self.start(query, options)
+      new StatementClient.new(faraday_client(options), query, options)
+    end
+
+    def self.resume(next_uri, options)
+      new StatementClient.new(faraday_client(options), nil, options, next_uri)
+    end
+
+    def self.faraday_client(options)
+      server = options[:server]
+      unless server
+        raise ArgumentError, ":server option is required"
+      end
+
+      ssl = faraday_ssl_options(options)
+
+      url = "#{ssl ? "https" : "http"}://#{server}"
+      proxy = options[:http_proxy] || options[:proxy]  # :proxy is obsoleted
+
+      faraday = Faraday.new(url: url, proxy: "#{proxy}", ssl: ssl) do |faraday|
+        #faraday.request :url_encoded
+        faraday.response :logger if options[:http_debug]
+        faraday.adapter Faraday.default_adapter
+      end
+
+      return faraday
+    end
+
+    def self.faraday_ssl_options(options)
+      ssl = options[:ssl]
+
+      case ssl
+      when true
+        ssl = {verify: true}
+
+      when Hash
+        verify = ssl.fetch(:verify, true)
+        case verify
+        when true
+          # detailed SSL options. pass through to faraday
+        when nil, false
+          ssl = {verify: false}
+        else
+          raise ArgumentError, "Can't convert #{verify.class} of :verify option of :ssl option to true or false"
+        end
+
+      when nil, false
+        ssl = false
+
+      else
+        raise ArgumentError, "Can't convert #{ssl.class} of :ssl option to true, false, or Hash"
+      end
+
+      return ssl
+    end
+
+    private_class_method :faraday_client, :faraday_ssl_options
+
+    def initialize(api)
+      @api = api
+    end
+
+    def current_results
+      @api.current_results
+    end
+
+    def advance
+      @api.advance
+    end
+
+    def wait_for_columns
+      while @api.current_results.columns == nil && @api.advance
+      end
+    end
+
+    def wait_for_data
+      while @api.current_results.data == nil && @api.advance
+      end
+    end
+
+    private :wait_for_columns
+    private :wait_for_data
+
+    def columns
+      wait_for_columns
+
+      raise_if_failed
+
+      return @api.current_results.columns
+    end
+
+    def rows
+      rows = []
+      each_row_chunk {|chunk|
+        rows.concat(chunk)
+      }
+      return rows
+    end
+
+    def each_row(&block)
+      each_row_chunk {|chunk|
+        chunk.each(&block)
+      }
+    end
+
+    def each_row_chunk(&block)
+      wait_for_data
+
+      raise_if_failed
+
+      if self.columns == nil
+        raise PrestoError, "Query #{@api.current_results.id} has no columns"
+      end
+
+      begin
+        if data = @api.current_results.data
+          block.call(data)
+        end
+      end while @api.advance
+    end
+
+    def query_info
+      @api.query_info
+    end
+
+    def next_uri
+      @api.current_results.next_uri
+    end
+
+    def cancel
+      @api.cancel_leaf_stage
+    end
+
+    def close
+      @api.cancel_leaf_stage
+      nil
+    end
+
+    def raise_if_failed
+      if @api.closed?
+        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
+        raise PrestoQueryError.new("Query #{results.id} failed: #{error.message}", results.id, error.error_code, error.failure_info)
+      end
+    end
+  end
+
+end

data/lib/presto/client/statement_client.rb

@@ -0,0 +1,259 @@
+#
+# 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 PrestoLegacy::Client
+
+  require 'multi_json'
+  require 'presto/client/models'
+  require 'presto/client/errors'
+
+  module PrestoHeaders
+    PRESTO_USER = "X-Presto-User"
+    PRESTO_SOURCE = "X-Presto-Source"
+    PRESTO_CATALOG = "X-Presto-Catalog"
+    PRESTO_SCHEMA = "X-Presto-Schema"
+    PRESTO_TIME_ZONE = "X-Presto-Time-Zone"
+    PRESTO_LANGUAGE = "X-Presto-Language"
+    PRESTO_SESSION = "X-Presto-Session"
+
+    PRESTO_CURRENT_STATE = "X-Presto-Current-State"
+    PRESTO_MAX_WAIT = "X-Presto-Max-Wait"
+    PRESTO_MAX_SIZE = "X-Presto-Max-Size"
+    PRESTO_PAGE_SEQUENCE_ID = "X-Presto-Page-Sequence-Id"
+  end
+
+  class StatementClient
+    HEADERS = {
+      "User-Agent" => "presto-ruby/#{VERSION}",
+    }
+
+    def initialize(faraday, query, options, next_uri=nil)
+      @faraday = faraday
+      @faraday.headers.merge!(HEADERS)
+
+      @options = options
+      @query = query
+      @closed = false
+      @exception = nil
+
+      @faraday.headers.merge!(optional_headers)
+
+      if next_uri
+        body = faraday_get_with_retry(next_uri)
+        @results = Models::QueryResults.decode(MultiJson.load(body))
+      else
+        post_query_request!
+      end
+    end
+
+    def optional_headers
+      headers = {}
+      if v = @options[:user]
+        headers[PrestoHeaders::PRESTO_USER] = v
+      end
+      if v = @options[:source]
+        headers[PrestoHeaders::PRESTO_SOURCE] = v
+      end
+      if v = @options[:catalog]
+        headers[PrestoHeaders::PRESTO_CATALOG] = v
+      end
+      if v = @options[:schema]
+        headers[PrestoHeaders::PRESTO_SCHEMA] = v
+      end
+      if v = @options[:time_zone]
+        headers[PrestoHeaders::PRESTO_TIME_ZONE] = v
+      end
+      if v = @options[:language]
+        headers[PrestoHeaders::PRESTO_LANGUAGE] = v
+      end
+      if v = @options[:properties]
+        headers[PrestoHeaders::PRESTO_SESSION] = encode_properties(v)
+      end
+      headers
+    end
+
+    private :optional_headers
+
+    def init_request(req)
+      req.options.timeout = @options[:http_timeout] || 300
+      req.options.open_timeout = @options[:http_open_timeout] || 60
+    end
+
+    private :init_request
+
+    def post_query_request!
+      uri = "/v1/statement"
+      response = @faraday.post do |req|
+        req.url uri
+
+        req.body = @query
+        init_request(req)
+      end
+
+      # TODO error handling
+      if response.status != 200
+        raise PrestoHttpError.new(response.status, "Failed to start query: #{response.body}")
+      end
+
+      body = response.body
+      @results = load_json(uri, body, Models::QueryResults)
+    end
+
+    private :post_query_request!
+
+    attr_reader :query
+
+    def debug?
+      !!@options[:debug]
+    end
+
+    def closed?
+      @closed
+    end
+
+    attr_reader :exception
+
+    def exception?
+      @exception
+    end
+
+    def query_failed?
+      @results.error != nil
+    end
+
+    def query_succeeded?
+      @results.error == nil && !@exception && !@closed
+    end
+
+    def current_results
+      @results
+    end
+
+    def has_next?
+      !!@results.next_uri
+    end
+
+    def advance
+      if closed? || !has_next?
+        return false
+      end
+      uri = @results.next_uri
+
+      body = faraday_get_with_retry(uri)
+      @results = load_json(uri, body, Models::QueryResults)
+
+      return true
+    end
+
+    def query_info
+      uri = "/v1/query/#{@results.id}"
+      body = faraday_get_with_retry(uri)
+      load_json(uri, body, Models::QueryInfo)
+    end
+
+    def load_json(uri, body, body_class)
+      hash = MultiJson.load(body)
+      begin
+        body_class.decode(hash)
+      rescue => e
+        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
+      end
+    end
+
+    private :load_json
+
+    def faraday_get_with_retry(uri, &block)
+      start = Time.now
+      attempts = 0
+
+      begin
+        begin
+          response = @faraday.get(uri)
+        rescue Faraday::Error::TimeoutError, Faraday::Error::ConnectionFailed
+          # temporally error to retry
+          response = nil
+        rescue => e
+          @exception = e
+          raise @exception
+        end
+
+        if response
+          if response.status == 200 && !response.body.to_s.empty?
+            return response.body
+          end
+
+          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
+          end
+        end
+
+        attempts += 1
+        sleep attempts * 0.1
+      end while (Time.now - start) < 2*60*60 && !@closed
+
+      @exception = PrestoHttpError.new(408, "Presto API error due to timeout")
+      raise @exception
+    end
+
+    def cancel_leaf_stage
+      if uri = @results.next_uri
+        response = @faraday.delete do |req|
+          req.url uri
+        end
+        return response.status / 100 == 2
+      end
+      return false
+    end
+
+    HTTP11_SEPARATOR = ["(", ")", "<", ">", "@", ",", ";", ":", "\\", "<", ">", "/", "[", "]", "?", "=", "{", "}", " ", "\v"]
+    HTTP11_TOKEN_CHARSET = (32..126).map {|x| x.chr } - HTTP11_SEPARATOR
+    HTTP11_TOKEN_REGEXP = /^[#{Regexp.escape(HTTP11_TOKEN_CHARSET.join)}]+\z/
+    HTTP11_CTL_CHARSET = (0..31).map {|x| x.chr } + [127.chr]
+    HTTP11_CTL_CHARSET_REGEXP = /[#{Regexp.escape(HTTP11_CTL_CHARSET.join)}]/
+
+    def encode_properties(properties)
+      # this is a hack to set same header multiple times.
+      properties.map do |k, v|
+        token = k.to_s
+        field_value = v.to_s  # TODO LWS encoding is not implemented
+        unless k =~ HTTP11_TOKEN_REGEXP
+          raise Faraday::ClientError, "Key of properties can't include HTTP/1.1 control characters or separators (#{HTTP11_SEPARATOR.map {|c| c =~ /\s/ ? c.dump : c }.join(' ')})"
+        end
+        if field_value =~ HTTP11_CTL_CHARSET_REGEXP
+          raise Faraday::ClientError, "Value of properties can't include HTTP/1.1 control characters"
+        end
+        "#{token}=#{field_value}"
+      end.join("\r\n#{PrestoHeaders::PRESTO_SESSION}: ")
+    end
+
+    def close
+      return if @closed
+
+      # cancel running statement
+      # TODO make async reqeust and ignore response?
+      cancel_leaf_stage
+
+      @closed = true
+      nil
+    end
+  end
+
+end

data/lib/presto/client/version.rb

@@ -0,0 +1,20 @@
+#
+# 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 PrestoLegacy
+  module Client
+    VERSION = "0.4.17"
+  end
+end

data/lib/presto/client.rb

@@ -0,0 +1,23 @@
+#
+# 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 PrestoLegacy
+  module Client
+
+    require 'presto/client/version'
+    require 'presto/client/client'
+
+  end
+end

data/lib/presto-client-legacy.rb

@@ -0,0 +1 @@
+require 'presto/client'

data/modelgen/modelgen.rb

@@ -0,0 +1,91 @@
+
+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 = source_dir
+
+predefined_simple_classes = %w[StageId TaskId ConnectorSession]
+predefined_models = %w[DistributionSnapshot PlanNode EquiJoinClause WriterTarget PageBufferInfo DeleteHandle]
+
+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]
+enum_types = %w[QueryState StageState TaskState QueueState PlanDistribution OutputPartitioning Step SortOrder BufferState NullPartitioning BlockedReason ParameterKind FunctionKind PartitionFunctionHandle]
+
+root_models = %w[QueryResults QueryInfo] + %w[
+OutputNode
+ProjectNode
+TableScanNode
+ValuesNode
+AggregationNode
+MarkDistinctNode
+FilterNode
+WindowNode
+RowNumberNode
+TopNRowNumberNode
+LimitNode
+DistinctLimitNode
+TopNNode
+SampleNode
+SortNode
+RemoteSourceNode
+JoinNode
+SemiJoinNode
+IndexJoinNode
+IndexSourceNode
+TableWriterNode
+DeleteNode
+MetadataDeleteNode
+TableFinishNode
+UnnestNode
+ExchangeNode
+UnionNode
+EnforceSingleRowNode
+] + %w[InsertTableHandle OutputTableHandle TableHandle]
+
+name_mapping = Hash[*%w[
+StatementStats StageStats ClientStageStats
+ClientStageStats StageStats ClientStageStats
+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
+].map.with_index { |v,i| i % 2 == 0 ? v : (source_path + "/" + v) }]
+
+analyzer = PrestoModels::ModelAnalyzer.new(
+  source_path,
+  skip_models: predefined_models + predefined_simple_classes + assume_primitive + enum_types,
+  path_mapping: path_mapping,
+  name_mapping: name_mapping
+)
+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 }
+