sequel-impala 1.0.1 → 1.1.0

data/lib/impala/thrift_patch.rb

@@ -0,0 +1,42 @@
+require 'socket'
+
+module Thrift
+  module KeepAlive
+    # We'll override #open so that once the socket is opened
+    # we enable keepalive on it
+    #
+    # Many queries are going to take a long time (10s of minutes) to complete
+    # and we don't want the connection to close while we wait for the
+    # query to return.
+    #
+    # Unfortunately, Thrift doesn't supply an easy way to get to the
+    # socket that it opens to communicate with Impala.
+    #
+    # I figured that while I was in here, monkey-patching a way to get
+    # to the socket, I might as well just enable keepalive here
+    # instead.
+    def open
+      super
+      s = @transport.handle
+      s.setsockopt(::Socket::SOL_SOCKET, ::Socket::SO_KEEPALIVE, true)
+
+      # Apparently Mac OS X (Darwin) doesn't implement the SOL_TCP options below
+      # so we'll hope keep alive works under Mac OS X, but in production
+      # we Dockerize Jigsaw, so these options should be available when
+      # we're running on Linux
+      if defined? ::Socket::SOL_TCP
+        s.setsockopt(::Socket::SOL_TCP, ::Socket::TCP_KEEPIDLE, 60)
+        s.setsockopt(::Socket::SOL_TCP, ::Socket::TCP_KEEPINTVL, 10)
+        s.setsockopt(::Socket::SOL_TCP, ::Socket::TCP_KEEPCNT, 5)
+      end
+    end
+  end
+
+  class BufferedTransport
+    prepend KeepAlive
+  end
+
+  class ImpalaSaslClientTransport
+    prepend KeepAlive
+  end
+end

data/lib/rbhive/connection.rb

@@ -19,18 +19,18 @@ module RBHive
     end
   end
   module_function :connect
-  
+
   class StdOutLogger
-    %w(fatal error warn info debug).each do |level| 
+    %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)
@@ -40,43 +40,43 @@ module RBHive
       @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)
@@ -85,7 +85,7 @@ module RBHive
         ResultSet.new(rows, the_schema)
       end
     end
-    
+
     def fetch_in_batch(query, batch_size=1_000)
       safe do
         execute_unsafe(query)
@@ -95,7 +95,7 @@ module RBHive
         end
       end
     end
-    
+
     def first(query)
       safe do
         execute_unsafe(query)
@@ -104,43 +104,43 @@ module RBHive
         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 }

data/lib/rbhive/explain_result.rb

@@ -2,33 +2,33 @@ 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|
@@ -43,4 +43,4 @@ class ExplainResult
       sections
     end
   end
-end
+end

data/lib/rbhive/schema_definition.rb

@@ -3,10 +3,10 @@ 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 = { 
+    TYPES = {
       :boolean  => :to_s,
       :string   => :to_s,
       :bigint   => :to_i,
@@ -16,16 +16,16 @@ module RBHive
       :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.
@@ -36,7 +36,7 @@ module RBHive
         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
@@ -46,23 +46,23 @@ module RBHive
         schema_names
       end
     end
-  
+
     def column_type_map
-      @column_type_map ||= column_names.inject({}) do |hsh, c| 
+      @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")
@@ -71,11 +71,11 @@ module RBHive
       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

data/lib/rbhive/t_c_l_i_connection.rb

@@ -30,7 +30,7 @@ module Thrift
 end
 
 module RBHive
-  
+
   HIVE_THRIFT_MAPPING = {
     10 => 0,
     11 => 1,
@@ -85,11 +85,11 @@ module RBHive
       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
@@ -97,7 +97,7 @@ module RBHive
       @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)
@@ -105,11 +105,11 @@ module RBHive
       @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]
@@ -188,7 +188,7 @@ module RBHive
       @logger.info("Setting #{name}=#{value}")
       self.execute("SET #{name}=#{value}")
     end
-    
+
     # Async execute
     def async_execute(query)
       @logger.info("Executing query asynchronously: #{query}")
@@ -204,35 +204,35 @@ module RBHive
 
       # Return handles to get hold of this query / session again
       {
-        session: @session.sessionHandle, 
-        guid: op_handle.operationId.guid, 
+        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(
@@ -262,18 +262,18 @@ module RBHive
         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)
@@ -290,7 +290,7 @@ module RBHive
         yield rows
       end
     end
-    
+
     def async_close_session(handles)
       validate_handles!(handles)
       @client.CloseSession(Hive2::Thrift::TCloseSessionReq.new( sessionHandle: handles[:session] ))
@@ -311,14 +311,16 @@ module RBHive
           vals = row.colVals
           cols.each do |i, col, conv|
             v = vals[i].get_value.value
-            h[col] = conv ? conv[v] : v
+            h[col] = unless v.nil?
+              conv ? conv[v] : v
+            end
           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)
@@ -327,7 +329,7 @@ module RBHive
       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)
@@ -346,7 +348,7 @@ module RBHive
 
       # Get search operation handle to fetch the results
       op_handle = exec_result.operationHandle
-      
+
       # Fetch the rows
       fetch_rows(op_handle, :first, max_rows)
     end
@@ -355,7 +357,7 @@ module RBHive
     # 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)
@@ -433,13 +435,13 @@ module RBHive
         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}"