sequel-impala 1.0.0 → 1.0.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 1583131d7fd60fe171fa38b12e51553760f4636f
-  data.tar.gz: 85c496bf9bdf022c29d9a309d7a1e2f2b202a288
+  metadata.gz: 03ea5f2607bc4908064302d49640df3a4e34eaa3
+  data.tar.gz: 18c93756bb5918f32cb6f32612856963f634e966
 SHA512:
-  metadata.gz: 82f28a92091bd3a46992f2eceac625ed30f9107cdc80e6495af4a45b76fcd55b4df23c0cac768ccce48038d3ba888e018db8df5184fc8846e7f17d793a0d6ff6
-  data.tar.gz: 93809720a5187f6c9eafe2708bf4cf86c1575caa5a0079077f7fef2d3b76a2dac6420b94c4b8b44d7a3e9a6336bdc049accce88ea80b301dee7049684c6ea135
+  metadata.gz: 72ca2b1c7177ecc8c2db06e8b266f4a1ff67085cf7ef77d464e6b7667ce896870d905a691acf4fc2b34e67c90777dc3767cc9b3a1e3a2a9252967beaf507b566
+  data.tar.gz: fb15a8bf19c03e54179666df88da7ffc0f0be90d48a65ada20722da03f46c4fa26d910a2f3c554a93149ea29f92f565fa3bbeb43b875c8ceacfcf413b64aa161

data/CHANGELOG

@@ -1,3 +1,19 @@
 === HEAD
 
+* Fix disconnect detection in impala and rbhive adapters (jeremyevans)
+
+* Make implicit qualify return an SQL::Identifier if given an unqualified string (jeremyevans)
+
+* Fix :search_path option handling when using Sequel::SQL::AliasedExpressions (jeremyevans)
+
+* Speed up multi_insert and import (jeremyevans)
+
+* Add rbhive adapter (jeremyevans)
+
+* Add :empty_null=>:ruby option to csv_to_parquet extension, which can support quoted CSV cells (jeremyevans)
+
+* Optimize csv_to_parquet extension by not spawning shells or unnecessary processes (jeremyevans)
+
+=== 1.0.0 (2015-12-04)
+
 * Initial Public Release

data/LICENSE

@@ -1,5 +1,6 @@
-Copyright (c) 2015 Jeremy Evans
+Copyright (c) 2015-2016 Jeremy Evans
 Copyright (c) 2013 Colin Marc
+Copyright (c) [2013] [Forward3D]
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to

data/README.md

@@ -0,0 +1,45 @@
+# sequel-impala
+
+sequel-impala adds support for Sequel to connect to the Impala database
+via the included impala driver, and the included jdbc-hive2 driver under JRuby.
+
+# Source Code
+
+Source code is available on GitHub at https://github.com/outcomesinsights/sequel-impala
+
+# Usage
+
+After installation, Sequel will automatically pick up the adapter as long as
+the lib directory is in RUBYLIB, if you use a connection string starting with
+`impala`, or `jdbc:hive2` on JRuby.
+
+# Connection Strings
+
+If using the impala driver (default host is localhost, default port is 21000):
+
+    impala://host:port
+
+If using the jdbc:hive2 driver on JRuby (port 21050 works in testing):
+
+    jdbc:hive2://host:port/;auth=noSasl
+
+# Dependencies
+
+* sequel 4+
+* thrift gem
+
+# License
+
+MIT/Apache
+
+# Author
+
+Ryan Duryea <aguynamedryan@gmail.com>
+
+Work on sequel-impala is generously funded by [Outcomes Insights, Inc.](http://outins.com)
+
+# Previous Author
+
+Jeremy Evans <code@jeremyevans.net>
+
+Provided initial work on this gem, and continues to maintain [Sequel](http://sequel.jeremyevans.net/).  We can't thank you enough!

data/lib/rbhive.rb

@@ -0,0 +1,8 @@
+require File.join(File.dirname(__FILE__), 'rbhive', 'connection')
+require File.join(File.dirname(__FILE__), 'rbhive', 'table_schema')
+require File.join(File.dirname(__FILE__), 'rbhive', 'result_set')
+require File.join(File.dirname(__FILE__), 'rbhive', 'explain_result')
+require File.join(File.dirname(__FILE__), 'rbhive', 'schema_definition')
+require File.join(File.dirname(__FILE__), *%w[rbhive t_c_l_i_result_set])
+require File.join(File.dirname(__FILE__), *%w[rbhive t_c_l_i_schema_definition])
+require File.join(File.dirname(__FILE__), *%w[rbhive t_c_l_i_connection])

data/lib/rbhive/connection.rb

@@ -0,0 +1,150 @@
+# suppress warnings
+old_verbose, $VERBOSE = $VERBOSE, nil
+# require thrift autogenerated files
+require File.join(File.split(File.dirname(__FILE__)).first, *%w[thrift thrift_hive])
+# require 'thrift'
+# restore warnings
+$VERBOSE = old_verbose
+
+module RBHive
+  def connect(server, port=10_000)
+    connection = RBHive::Connection.new(server, port)
+    ret = nil
+    begin
+      connection.open
+      ret = yield(connection)
+    ensure
+      connection.close
+      ret
+    end
+  end
+  module_function :connect
+  
+  class StdOutLogger
+    %w(fatal error warn info debug).each do |level| 
+      define_method level.to_sym do |message|
+        STDOUT.puts(message)
+     end
+   end
+  end
+  
+  class Connection
+    attr_reader :client
+    
+    def initialize(server, port=10_000, logger=StdOutLogger.new)
+      @socket = Thrift::Socket.new(server, port)
+      @transport = Thrift::BufferedTransport.new(@socket)
+      @protocol = Thrift::BinaryProtocol.new(@transport)
+      @client = Hive::Thrift::ThriftHive::Client.new(@protocol)
+      @logger = logger
+      @logger.info("Connecting to #{server} on port #{port}")
+      @mutex = Mutex.new
+    end
+    
+    def open
+      @transport.open
+    end
+    
+    def close
+      @transport.close
+    end
+    
+    def client
+      @client
+    end
+    
+    def execute(query)
+      execute_safe(query)
+    end
+    
+    def explain(query)
+      safe do
+        execute_unsafe("EXPLAIN "+ query)
+        ExplainResult.new(client.fetchAll)
+      end
+    end
+    
+    def priority=(priority)
+      set("mapred.job.priority", priority)
+    end
+    
+    def queue=(queue)
+      set("mapred.job.queue.name", queue)
+    end
+    
+    def set(name,value)
+      @logger.info("Setting #{name}=#{value}")
+      client.execute("SET #{name}=#{value}")
+    end
+    
+    def fetch(query)
+      safe do
+        execute_unsafe(query)
+        rows = client.fetchAll
+        the_schema = SchemaDefinition.new(client.getSchema, rows.first)
+        ResultSet.new(rows, the_schema)
+      end
+    end
+    
+    def fetch_in_batch(query, batch_size=1_000)
+      safe do
+        execute_unsafe(query)
+        until (next_batch = client.fetchN(batch_size)).empty?
+          the_schema ||= SchemaDefinition.new(client.getSchema, next_batch.first)
+          yield ResultSet.new(next_batch, the_schema)
+        end
+      end
+    end
+    
+    def first(query)
+      safe do
+        execute_unsafe(query)
+        row = client.fetchOne
+        the_schema = SchemaDefinition.new(client.getSchema, row)
+        ResultSet.new([row], the_schema).first
+      end
+    end
+    
+    def schema(example_row=[])
+      safe { SchemaDefinition.new(client.getSchema, example_row) }
+    end
+    
+    def create_table(schema)
+      execute(schema.create_table_statement)
+    end
+    
+    def drop_table(name)
+      name = name.name if name.is_a?(TableSchema)
+      execute("DROP TABLE `#{name}`")
+    end
+    
+    def replace_columns(schema)
+      execute(schema.replace_columns_statement)
+    end
+    
+    def add_columns(schema)
+      execute(schema.add_columns_statement)
+    end
+    
+    def method_missing(meth, *args)
+      client.send(meth, *args)
+    end
+    
+    private
+    
+    def execute_safe(query)
+      safe { execute_unsafe(query) }
+    end
+    
+    def execute_unsafe(query)
+      @logger.info("Executing Hive Query: #{query}")
+      client.execute(query)
+    end
+    
+    def safe
+      ret = nil
+      @mutex.synchronize { ret = yield }
+      ret
+    end
+  end
+end

data/lib/rbhive/explain_result.rb

@@ -0,0 +1,46 @@
+class ExplainResult
+  def initialize(rows)
+    @rows = rows
+  end
+  
+  def ast
+    by_section[:abstract_syntax_tree].first
+  end
+  
+  def stage_count
+    stage_dependencies.length
+  end
+  
+  def stage_dependencies
+    by_section[:stage_dependencies] || []
+  end
+  
+  def to_tsv
+    @rows.join("\n")
+  end
+  
+  def raw
+    @rows
+  end
+  
+  def to_s
+    to_tsv
+  end
+  
+  private
+  
+  def by_section
+    current_section = nil
+    @rows.inject({}) do |sections, row|
+      if row.match(/^[A-Z]/)
+        current_section = row.chomp(':').downcase.gsub(' ', '_').to_sym
+        sections[current_section] = []
+      elsif row.length == 0
+        next sections
+      else
+        sections[current_section] << row.strip
+      end
+      sections
+    end
+  end
+end

data/lib/rbhive/result_set.rb

@@ -0,0 +1,37 @@
+module RBHive
+  class ResultSet < Array
+    def initialize(rows, schema)
+      @schema = schema
+      super(rows.map {|r| @schema.coerce_row(r) })
+    end
+
+    def column_names
+      @schema.column_names
+    end
+
+    def column_type_map
+      @schema.column_type_map
+    end
+
+    def to_csv(out_file=nil)
+      to_separated_output(",", out_file)
+    end
+
+    def to_tsv(out_file=nil)
+      to_separated_output("\t", out_file)
+    end
+
+    def as_arrays
+      @as_arrays ||= self.map{ |r| @schema.coerce_row_to_array(r) }
+    end
+
+    private
+
+    def to_separated_output(sep, out_file)
+      rows = self.map { |r| @schema.coerce_row_to_array(r).join(sep) }
+      sv = rows.join("\n")
+      return sv if out_file.nil?
+      File.open(out_file, 'w+') { |f| f << sv }
+    end
+  end
+end

data/lib/rbhive/schema_definition.rb

@@ -0,0 +1,86 @@
+require 'json'
+
+module RBHive
+  class SchemaDefinition
+    attr_reader :schema
+  
+    NAN = Float::NAN rescue 0.0/0.0
+    INFINITY = Float::INFINITY rescue 1.0/0.0
+    TYPES = { 
+      :boolean  => :to_s,
+      :string   => :to_s,
+      :bigint   => :to_i,
+      :float    => :to_f,
+      :double   => :to_f,
+      :int      => :to_i,
+      :smallint => :to_i,
+      :tinyint  => :to_i,
+    }
+  
+    def initialize(schema, example_row)
+      @schema = schema
+      @example_row = example_row ? example_row.split("\t") : []
+    end
+  
+    def column_names
+      @column_names ||= begin
+        schema_names = @schema.fieldSchemas.map {|c| c.name }
+        
+        # In rare cases Hive can return two identical column names
+        # consider SELECT a.foo, b.foo...
+        # in this case you get two columns called foo with no disambiguation.
+        # as a (far from ideal) solution we detect this edge case and rename them
+        # a.foo => foo1, b.foo => foo2
+        # otherwise we will trample one of the columns during Hash mapping.
+        s = Hash.new(0)
+        schema_names.map! { |c| s[c] += 1; s[c] > 1 ? "#{c}---|---#{s[c]}" : c }
+        schema_names.map! { |c| s[c] > 1 ? "#{c}---|---1" : c }
+        schema_names.map! { |c| c.gsub('---|---', '_').to_sym }
+        
+        # Lets fix the fact that Hive doesn't return schema data for partitions on SELECT * queries
+        # For now we will call them :_p1, :_p2, etc. to avoid collisions.
+        offset = 0
+        while schema_names.length < @example_row.length
+          schema_names.push(:"_p#{offset+=1}")
+        end
+        schema_names
+      end
+    end
+  
+    def column_type_map
+      @column_type_map ||= column_names.inject({}) do |hsh, c| 
+        definition = @schema.fieldSchemas.find {|s| s.name.to_sym == c }
+        # If the column isn't in the schema (eg partitions in SELECT * queries) assume they are strings
+        hsh[c] = definition ? definition.type.to_sym : :string
+        hsh
+      end
+    end
+  
+    def coerce_row(row)
+      column_names.zip(row.split("\t")).inject({}) do |hsh, (column_name, value)|
+        hsh[column_name] = coerce_column(column_name, value)
+        hsh
+      end
+    end
+  
+    def coerce_column(column_name, value)
+      type = column_type_map[column_name]
+      return INFINITY if (type != :string && value == "Infinity")
+      return NAN if (type != :string && value == "NaN")
+      return coerce_complex_value(value) if type.to_s =~ /^array/
+      conversion_method = TYPES[type]
+      conversion_method ? value.send(conversion_method) : value
+    end
+  
+    def coerce_row_to_array(row)
+      column_names.map { |n| row[n] }
+    end
+    
+    def coerce_complex_value(value)
+      return nil if value.nil?
+      return nil if value.length == 0
+      return nil if value == 'null'
+      JSON.parse(value)
+    end
+  end
+end

data/lib/rbhive/t_c_l_i_connection.rb

@@ -0,0 +1,464 @@
+# suppress warnings
+old_verbose, $VERBOSE = $VERBOSE, nil
+
+raise 'Thrift is not loaded' unless defined?(Thrift)
+raise 'RBHive is not loaded' unless defined?(RBHive)
+
+# require thrift autogenerated files
+require File.join(File.dirname(__FILE__), *%w[.. thrift t_c_l_i_service_constants])
+require File.join(File.dirname(__FILE__), *%w[.. thrift t_c_l_i_service])
+require File.join(File.dirname(__FILE__), *%w[.. thrift sasl_client_transport])
+
+# restore warnings
+$VERBOSE = old_verbose
+
+# Monkey patch thrift to set an infinite read timeout
+module Thrift
+  class HTTPClientTransport < BaseTransport
+    def flush
+      http = Net::HTTP.new @url.host, @url.port
+      http.use_ssl = @url.scheme == 'https'
+      http.read_timeout = nil
+      http.verify_mode = @ssl_verify_mode if @url.scheme == 'https'
+      resp = http.post(@url.request_uri, @outbuf, @headers)
+      data = resp.body
+      data = Bytes.force_binary_encoding(data)
+      @inbuf = StringIO.new data
+      @outbuf = Bytes.empty_byte_buffer
+    end
+  end
+end
+
+module RBHive
+  
+  HIVE_THRIFT_MAPPING = {
+    10 => 0,
+    11 => 1,
+    12 => 2,
+    13 => 6,
+    :cdh4 => 0,
+    :cdh5 => 4,
+    :PROTOCOL_V1 => 0,
+    :PROTOCOL_V2 => 1,
+    :PROTOCOL_V3 => 2,
+    :PROTOCOL_V4 => 3,
+    :PROTOCOL_V5 => 4,
+    :PROTOCOL_V6 => 5,
+    :PROTOCOL_V7 => 6
+  }
+
+  def tcli_connect(server, port = 10_000, options={})
+    logger = options.key?(:logger) ? options.delete(:logger) : StdOutLogger.new
+    connection = RBHive::TCLIConnection.new(server, port, options, logger)
+    ret = nil
+    begin
+      connection.open
+      connection.open_session
+      ret = yield(connection)
+
+    ensure
+      # Try to close the session and our connection if those are still open, ignore io errors
+      begin
+        connection.close_session if connection.session
+        connection.close
+      rescue IOError => e
+        # noop
+      end
+    end
+
+    ret
+  end
+  module_function :tcli_connect
+
+  class StdOutLogger
+    %w(fatal error warn info debug).each do |level|
+      define_method level.to_sym do |message|
+        STDOUT.puts(message)
+     end
+   end
+  end
+
+  class TCLIConnection
+    attr_reader :client
+
+    def initialize(server, port = 10_000, options = {}, logger = StdOutLogger.new)
+      options ||= {} # backwards compatibility
+      raise "'options' parameter must be a hash" unless options.is_a?(Hash)
+      @sasl_params = options.delete(:sasl_params) || {}
+      
+      if options[:transport] == :sasl and @sasl_params.empty?
+        raise ":transport is set to :sasl, but no :sasl_params option was supplied"
+      end
+      
+      # Defaults to buffered transport, Hive 0.10, 1800 second timeout
+      options[:transport]     ||= :buffered
+      options[:hive_version]  ||= 10
+      options[:timeout]       ||= 1800
+      @options = options
+      # Look up the appropriate Thrift protocol version for the supplied Hive version
+      @thrift_protocol_version = thrift_hive_protocol(options[:hive_version])
+      
+      @logger = logger
+      @transport = thrift_transport(server, port)
+      @protocol = Thrift::BinaryProtocol.new(@transport)
+      @client = Hive2::Thrift::TCLIService::Client.new(@protocol)
+      @session = nil
+      @logger.info("Connecting to HiveServer2 #{server} on port #{port}")
+    end
+    
+    def thrift_hive_protocol(version)
+      HIVE_THRIFT_MAPPING[version] || raise("Invalid Hive version")
+    end
+    
+    def thrift_transport(server, port)
+      @logger.info("Initializing transport #{@options[:transport]}")
+      case @options[:transport]
+      when :buffered
+        return Thrift::BufferedTransport.new(thrift_socket(server, port, @options[:timeout]))
+      when :sasl
+        return Thrift::SaslClientTransport.new(thrift_socket(server, port, @options[:timeout]),
+                                               parse_sasl_params(@sasl_params))
+      when :http
+        return Thrift::HTTPClientTransport.new("http://#{server}:#{port}/cliservice")
+      else
+        raise "Unrecognised transport type '#{transport}'"
+      end
+    end
+
+    def thrift_socket(server, port, timeout)
+      socket = Thrift::Socket.new(server, port)
+      socket.timeout = timeout
+      socket
+    end
+
+    # Processes SASL connection params and returns a hash with symbol keys or a nil
+    def parse_sasl_params(sasl_params)
+      # Symbilize keys in a hash
+      if sasl_params.kind_of?(Hash)
+        return sasl_params.inject({}) do |memo,(k,v)|
+          memo[k.to_sym] = v;
+          memo
+        end
+      end
+      return nil
+    end
+
+    def open
+      @transport.open
+    end
+
+    def close
+      @transport.close
+    end
+
+    def open_session
+      @session = @client.OpenSession(prepare_open_session(@thrift_protocol_version))
+    end
+
+    def close_session
+      @client.CloseSession prepare_close_session
+      @session = nil
+    end
+
+    def session
+      @session && @session.sessionHandle
+    end
+
+    def client
+      @client
+    end
+
+    def execute(query)
+      @logger.info("Executing Hive Query: #{query}")
+      req = prepare_execute_statement(query)
+      exec_result = client.ExecuteStatement(req)
+      raise_error_if_failed!(exec_result)
+      exec_result
+    end
+
+    def priority=(priority)
+      set("mapred.job.priority", priority)
+    end
+
+    def queue=(queue)
+      set("mapred.job.queue.name", queue)
+    end
+
+    def set(name,value)
+      @logger.info("Setting #{name}=#{value}")
+      self.execute("SET #{name}=#{value}")
+    end
+    
+    # Async execute
+    def async_execute(query)
+      @logger.info("Executing query asynchronously: #{query}")
+      exec_result = @client.ExecuteStatement(
+        Hive2::Thrift::TExecuteStatementReq.new(
+          sessionHandle: @session.sessionHandle,
+          statement: query,
+          runAsync: true
+        )
+      )
+      raise_error_if_failed!(exec_result)
+      op_handle = exec_result.operationHandle
+
+      # Return handles to get hold of this query / session again
+      {
+        session: @session.sessionHandle, 
+        guid: op_handle.operationId.guid, 
+        secret: op_handle.operationId.secret
+      }
+    end
+    
+    # Is the query complete?
+    def async_is_complete?(handles)
+      async_state(handles) == :finished
+    end
+    
+    # Is the query actually running?
+    def async_is_running?(handles)
+      async_state(handles) == :running
+    end
+    
+    # Has the query failed?
+    def async_is_failed?(handles)
+      async_state(handles) == :error
+    end
+    
+    def async_is_cancelled?(handles)
+      async_state(handles) == :cancelled
+    end
+    
+    def async_cancel(handles)
+      @client.CancelOperation(prepare_cancel_request(handles))
+    end
+    
+    # Map states to symbols
+    def async_state(handles)
+      response = @client.GetOperationStatus(
+        Hive2::Thrift::TGetOperationStatusReq.new(operationHandle: prepare_operation_handle(handles))
+      )
+
+      case response.operationState
+      when Hive2::Thrift::TOperationState::FINISHED_STATE
+        return :finished
+      when Hive2::Thrift::TOperationState::INITIALIZED_STATE
+        return :initialized
+      when Hive2::Thrift::TOperationState::RUNNING_STATE
+        return :running
+      when Hive2::Thrift::TOperationState::CANCELED_STATE
+        return :cancelled
+      when Hive2::Thrift::TOperationState::CLOSED_STATE
+        return :closed
+      when Hive2::Thrift::TOperationState::ERROR_STATE
+        return :error
+      when Hive2::Thrift::TOperationState::UKNOWN_STATE
+        return :unknown
+      when Hive2::Thrift::TOperationState::PENDING_STATE
+        return :pending
+      when nil
+        raise "No operation state found for handles - has the session been closed?"
+      else
+        return :state_not_in_protocol
+      end
+    end
+        
+    # Async fetch results from an async execute
+    def async_fetch(handles, max_rows = 100)
+      # Can't get data from an unfinished query
+      unless async_is_complete?(handles)
+        raise "Can't perform fetch on a query in state: #{async_state(handles)}"
+      end
+      
+      # Fetch and
+      fetch_rows(prepare_operation_handle(handles), :first, max_rows)
+    end
+    
+    # Performs a query on the server, fetches the results in batches of *batch_size* rows
+    # and yields the result batches to a given block as arrays of rows.
+    def async_fetch_in_batch(handles, batch_size = 1000, &block)
+      raise "No block given for the batch fetch request!" unless block_given?
+      # Can't get data from an unfinished query
+      unless async_is_complete?(handles)
+        raise "Can't perform fetch on a query in state: #{async_state(handles)}"
+      end
+
+      # Now let's iterate over the results
+      loop do
+        rows = fetch_rows(prepare_operation_handle(handles), :next, batch_size)
+        break if rows.empty?
+        yield rows
+      end
+    end
+    
+    def async_close_session(handles)
+      validate_handles!(handles)
+      @client.CloseSession(Hive2::Thrift::TCloseSessionReq.new( sessionHandle: handles[:session] ))
+    end
+
+    def get_column_info(op_handle)
+      cols = get_schema_for(op_handle).columns
+      [cols.map(&:columnName), cols.map{|c| c.typeDesc.types.first.primitiveEntry.type}]
+    end
+
+    def yield_hash_rows(op_handle, columns, convertors)
+      i = -1
+      cols = columns.zip(convertors).map{|col, conv| [i+=1, col, conv]}
+      rows = fetch_rows(op_handle)
+      until rows.empty?
+        rows.each do |row|
+          h = {}
+          vals = row.colVals
+          cols.each do |i, col, conv|
+            v = vals[i].get_value.value
+            h[col] = conv ? conv[v] : v
+          end
+          yield h
+        end
+        rows = fetch_rows(op_handle, :next)
+      end
+    end
+    
+    # Pull rows from the query result
+    def fetch_rows(op_handle, orientation = :first, max_rows = 1000)
+      fetch_req = prepare_fetch_results(op_handle, orientation, max_rows)
+      fetch_results = @client.FetchResults(fetch_req)
+      raise_error_if_failed!(fetch_results)
+      fetch_results.results.rows
+      #TCLIResultSet.new(rows, TCLISchemaDefinition.new(get_schema_for(op_handle), rows.first))
+    end
+        
+    # Performs a explain on the supplied query on the server, returns it as a ExplainResult.
+    # (Only works on 0.12 if you have this patch - https://issues.apache.org/jira/browse/HIVE-5492)
+    def explain(query)
+      rows = []
+      fetch_in_batch("EXPLAIN " + query) do |batch|
+        rows << batch.map { |b| b[:Explain] }
+      end
+      ExplainResult.new(rows.flatten)
+    end
+
+    # Performs a query on the server, fetches up to *max_rows* rows and returns them as an array.
+    def fetch(query, max_rows = 100)
+      # Execute the query and check the result
+      exec_result = execute(query)
+      raise_error_if_failed!(exec_result)
+
+      # Get search operation handle to fetch the results
+      op_handle = exec_result.operationHandle
+      
+      # Fetch the rows
+      fetch_rows(op_handle, :first, max_rows)
+    end
+
+    # Performs a query on the server, fetches the results in batches of *batch_size* rows
+    # and yields the result batches to a given block as arrays of rows.
+    def fetch_in_batch(query, batch_size = 1000, &block)
+      raise "No block given for the batch fetch request!" unless block_given?
+      
+      # Execute the query and check the result
+      exec_result = execute(query)
+      raise_error_if_failed!(exec_result)
+
+      # Get search operation handle to fetch the results
+      op_handle = exec_result.operationHandle
+
+      # Prepare fetch results request
+      fetch_req = prepare_fetch_results(op_handle, :next, batch_size)
+
+      # Now let's iterate over the results
+      loop do
+        rows = fetch_rows(op_handle, :next, batch_size)
+        break if rows.empty?
+        yield rows
+      end
+    end
+
+    def create_table(schema)
+      execute(schema.create_table_statement)
+    end
+
+    def drop_table(name)
+      name = name.name if name.is_a?(TableSchema)
+      execute("DROP TABLE `#{name}`")
+    end
+
+    def replace_columns(schema)
+      execute(schema.replace_columns_statement)
+    end
+
+    def add_columns(schema)
+      execute(schema.add_columns_statement)
+    end
+
+    def method_missing(meth, *args)
+      client.send(meth, *args)
+    end
+
+    private
+
+    def prepare_open_session(client_protocol)
+      req = ::Hive2::Thrift::TOpenSessionReq.new( @sasl_params.empty? ? [] : @sasl_params )
+      req.client_protocol = client_protocol
+      req
+    end
+
+    def prepare_close_session
+      ::Hive2::Thrift::TCloseSessionReq.new( sessionHandle: self.session )
+    end
+
+    def prepare_execute_statement(query)
+      ::Hive2::Thrift::TExecuteStatementReq.new( sessionHandle: self.session, statement: query.to_s, confOverlay: {"impala.resultset.cache.size"=>"100000"} )
+    end
+
+    def prepare_fetch_results(handle, orientation=:first, rows=100)
+      orientation_value = "FETCH_#{orientation.to_s.upcase}"
+      valid_orientations = ::Hive2::Thrift::TFetchOrientation::VALUE_MAP.values
+      unless valid_orientations.include?(orientation_value)
+        raise ArgumentError, "Invalid orientation: #{orientation.inspect}"
+      end
+      orientation_const = eval("::Hive2::Thrift::TFetchOrientation::#{orientation_value}")
+      ::Hive2::Thrift::TFetchResultsReq.new(
+        operationHandle: handle,
+        orientation: orientation_const,
+        maxRows: rows
+      )
+    end
+
+    def prepare_operation_handle(handles)
+      validate_handles!(handles)
+      Hive2::Thrift::TOperationHandle.new(
+        operationId: Hive2::Thrift::THandleIdentifier.new(guid: handles[:guid], secret: handles[:secret]),
+        operationType: Hive2::Thrift::TOperationType::EXECUTE_STATEMENT,
+        hasResultSet: false
+      )
+    end
+    
+    def prepare_cancel_request(handles)
+      Hive2::Thrift::TCancelOperationReq.new(
+        operationHandle: prepare_operation_handle(handles)
+      )
+    end
+    
+    def validate_handles!(handles)
+      unless handles.has_key?(:guid) and handles.has_key?(:secret) and handles.has_key?(:session)
+        raise "Invalid handles hash: #{handles.inspect}"
+      end
+    end
+
+    def get_schema_for(handle)
+      req = ::Hive2::Thrift::TGetResultSetMetadataReq.new( operationHandle: handle )
+      metadata = client.GetResultSetMetadata( req )
+      metadata.schema
+    end
+
+    # Raises an exception if given operation result is a failure
+    def raise_error_if_failed!(result)
+      return if result.status.statusCode == 0
+      error_message = result.status.errorMessage || 'Execution failed!'
+      raise RBHive::TCLIConnectionError.new(error_message)
+    end
+  end
+
+  class TCLIConnectionError < StandardError; end
+end