sequel-impala 1.0.0

data/README.rdoc

@@ -0,0 +1,39 @@
+= 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/jeremyevans/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 <tt>jdbc:hive2</tt> 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
+
+Jeremy Evans <code@jeremyevans.net>
+
+Work on sequel-impala is generously funded by Outcomes Insights, Inc.

data/Rakefile

@@ -0,0 +1,39 @@
+require "rake"
+require "rake/clean"
+
+CLEAN.include ["sequel-impala-*.gem", "rdoc"]
+
+desc "Build sequel-impala gem"
+task :package=>[:clean] do |p|
+  sh %{#{FileUtils::RUBY} -S gem build sequel-impala.gemspec}
+end
+
+### Specs
+
+desc "Run specs"
+task "spec" do
+  sh "#{FileUtils::RUBY} -rubygems -I lib -e 'ARGV.each{|f| require f}' ./spec/*_test.rb"
+end
+
+task :default => :spec
+
+### RDoc
+
+RDOC_DEFAULT_OPTS = ["--quiet", "--line-numbers", "--inline-source", '--title', 'sequel-impala: Sequel support for Impala database']
+
+begin
+  gem 'rdoc'
+  gem 'hanna-nouveau'
+  RDOC_DEFAULT_OPTS.concat(['-f', 'hanna'])
+rescue Gem::LoadError
+end
+
+RDOC_OPTS = RDOC_DEFAULT_OPTS + ['--main', 'README.rdoc']
+
+require 'rdoc/task'
+RDoc::Task.new do |rdoc|
+  rdoc.rdoc_dir = "rdoc"
+  rdoc.options += RDOC_OPTS
+  rdoc.rdoc_files.add %w"README.rdoc CHANGELOG LICENSE lib/**/*.rb"
+end
+

data/lib/impala.rb

@@ -0,0 +1,47 @@
+
+# the generated ruby files use a relative require, so we need to add the
+# generated directory to $LOAD_PATH
+this_dir = File.expand_path(File.dirname(__FILE__))
+gen_dir = File.join(this_dir, 'impala/protocol')
+$LOAD_PATH.push(gen_dir) unless $LOAD_PATH.include?(gen_dir)
+
+require 'impala/version'
+
+require 'thrift'
+require 'time'
+require 'impala/protocol'
+require 'impala/cursor'
+require 'impala/connection'
+
+module Impala
+  DEFAULT_HOST = 'localhost'
+  DEFAULT_PORT = 21000
+  class Error < StandardError; end
+  class InvalidQueryError < Error; end
+  class ConnectionError < Error; end
+  class CursorError < Error; end
+
+  # Connect to an Impala server. If a block is given, it will close the
+  # connection after yielding the connection to the block.
+  # @param [String] host the hostname or IP address of the Impala server
+  # @param [int] port the port that the Impala server is listening on
+  # @yieldparam [Connection] conn the open connection. Will be closed once the block
+  #    finishes
+  # @return [Connection] the open connection, or, if a block is
+  #    passed, the return value of the block
+  def self.connect(host=DEFAULT_HOST, port=DEFAULT_PORT)
+    connection = Connection.new(host, port)
+
+    if block_given?
+      begin
+        ret = yield connection
+      ensure
+        connection.close
+      end
+    else
+      ret = connection
+    end
+
+    ret
+  end
+end

data/lib/impala/connection.rb

@@ -0,0 +1,117 @@
+module Impala
+  # This object represents a connection to an Impala server. It can be used to
+  # perform queries on the database.
+  class Connection
+    LOG_CONTEXT_ID = "impala-ruby"
+
+    # Don't instantiate Connections directly; instead, use {Impala.connect}.
+    def initialize(host, port)
+      @host = host
+      @port = port
+      @connected = false
+      open
+    end
+
+    def inspect
+      "#<#{self.class} #{@host}:#{@port}#{open? ? '' : ' (DISCONNECTED)'}>"
+    end
+
+    # Open the connection if it's currently closed.
+    def open
+      return if @connected
+
+      socket = Thrift::Socket.new(@host, @port)
+
+      @transport = Thrift::BufferedTransport.new(socket)
+      @transport.open
+
+      proto = Thrift::BinaryProtocol.new(@transport)
+      @service = Protocol::ImpalaService::Client.new(proto)
+      @connected = true
+    end
+
+    # Close this connection. It can still be reopened with {#open}.
+    def close
+      return unless @connected
+
+      @transport.close
+      @connected = false
+    end
+
+    # Returns true if the connection is currently open.
+    def open?
+      @connected
+    end
+
+    # Refresh the metadata store.
+    def refresh
+      raise ConnectionError.new("Connection closed") unless open?
+      @service.ResetCatalog
+    end
+
+    # Perform a query and return all the results. This will
+    # load the entire result set into memory, so if you're dealing with lots
+    # of rows, {#execute} may work better.
+    # @param [String] query the query you want to run
+    # @param [Hash] query_options the options to set user and configuration
+    #   except for :user, see TImpalaQueryOptions in ImpalaService.thrift
+    # @option query_options [String] :user the user runs the query
+    # @return [Array<Hash>] an array of hashes, one for each row.
+    def query(raw_query, query_options = {})
+      execute(raw_query, query_options).fetch_all
+    end
+
+    # Perform a query and return a cursor for iterating over the results.
+    # @param [String] query the query you want to run
+    # @param [Hash] query_options the options to set user and configuration
+    #   except for :user, see TImpalaQueryOptions in ImpalaService.thrift
+    # @option query_options [String] :user the user runs the query
+    # @return [Cursor] a cursor for the result rows
+    def execute(raw_query, query_options = {})
+      raise ConnectionError.new("Connection closed") unless open?
+
+      query = sanitize_query(raw_query)
+      handle = send_query(query, query_options)
+
+      check_result(handle)
+      Cursor.new(handle, @service)
+    end
+
+    private
+
+    def sanitize_query(raw_query)
+      words = raw_query.split
+      raise InvalidQueryError.new("Empty query") if words.empty?
+
+      command = words.first.downcase
+      ([command] + words[1..-1]).join(' ')
+    end
+
+    def send_query(sanitized_query, query_options)
+      query = Protocol::Beeswax::Query.new
+      query.query = sanitized_query
+
+      query.hadoop_user = query_options.delete(:user) if query_options[:user]
+      query.configuration = query_options.map do |key, value|
+        "#{key.upcase}=#{value}"
+      end
+
+      @service.executeAndWait(query, LOG_CONTEXT_ID)
+    end
+
+    def check_result(handle)
+      state = @service.get_state(handle)
+      if state == Protocol::Beeswax::QueryState::EXCEPTION
+        close_handle(handle)
+        raise ConnectionError.new("The query was aborted")
+      end
+    rescue
+      close_handle(handle)
+      raise
+    end
+
+    def close_handle(handle)
+      @service.close(handle)
+    end
+  end
+end

data/lib/impala/cursor.rb

@@ -0,0 +1,157 @@
+module Impala
+  # Cursors are used to iterate over result sets without loading them all
+  # into memory at once. This can be useful if you're dealing with lots of
+  # rows. It implements Enumerable, so you can use each/select/map/etc.
+  class Cursor
+    BUFFER_SIZE = 1024
+    include Enumerable
+
+    def self.typecast_boolean(value)
+      value == 'true'
+    end
+
+    def self.typecast_int(value)
+      value.to_i
+    end
+
+    def self.typecast_float(value)
+      value.to_f
+    end
+
+    def self.typecast_decimal(value)
+      BigDecimal.new(value)
+    end
+
+    def self.typecast_timestamp(value)
+      Time.parse(value)
+    end
+
+    TYPECAST_MAP = {
+      'boolean'=>method(:typecast_boolean),
+      'int'=>method(:typecast_int),
+      'double'=>method(:typecast_float),
+      'decimal'=>method(:typecast_decimal),
+      'timestamp'=>method(:typecast_timestamp),
+    }
+    TYPECAST_MAP['tinyint'] = TYPECAST_MAP['smallint'] = TYPECAST_MAP['bigint'] = TYPECAST_MAP['int']
+    TYPECAST_MAP['float'] = TYPECAST_MAP['double']
+    TYPECAST_MAP.freeze
+
+    NULL = 'NULL'.freeze
+
+    attr_reader :columns
+
+    attr_reader :typecast_map
+
+    def initialize(handle, service)
+      @handle = handle
+      @service = service
+
+
+      @row_buffer = []
+      @done = false
+      @open = true
+      @typecast_map = TYPECAST_MAP.dup
+      @columns = metadata.schema.fieldSchemas.map(&:name)
+    end
+
+    def inspect
+      "#<#{self.class}#{open? ? '' : ' (CLOSED)'}>"
+    end
+
+    def each
+      while row = fetch_row
+        yield row
+      end
+    end
+
+    # Returns the next available row as a hash, or nil if there are none left.
+    # @return [Hash, nil] the next available row, or nil if there are none
+    #    left
+    # @see #fetch_all
+    def fetch_row
+      if @row_buffer.empty?
+        if @done
+          return nil
+        else
+          fetch_more
+        end
+      end
+
+      @row_buffer.shift
+    end
+
+    # Returns all the remaining rows in the result set.
+    # @return [Array<Hash>] the remaining rows in the result set
+    # @see #fetch_one
+    def fetch_all
+      self.to_a
+    end
+
+    # Close the cursor on the remote server. Once a cursor is closed, you
+    # can no longer fetch any rows from it.
+    def close
+      @open = false
+      @service.close(@handle)
+    end
+
+    # Returns true if the cursor is still open.
+    def open?
+      @open
+    end
+
+    # Returns true if there are any more rows to fetch.
+    def has_more?
+      !@done || !@row_buffer.empty?
+    end
+
+    def runtime_profile
+      @service.GetRuntimeProfile(@handle)
+    end
+
+    private
+
+    def metadata
+      @metadata ||= @service.get_results_metadata(@handle)
+    end
+
+    def fetch_more
+      fetch_batch until @done || @row_buffer.count >= BUFFER_SIZE
+    end
+
+    def fetch_batch
+      raise CursorError.new("Cursor has expired or been closed") unless @open
+
+      begin
+        res = @service.fetch(@handle, false, BUFFER_SIZE)
+      rescue Protocol::Beeswax::BeeswaxException
+        @open = false
+        raise CursorError.new("Cursor has expired or been closed")
+      end
+
+      rows = res.data.map { |raw| parse_row(raw) }
+      @row_buffer.concat(rows)
+
+      unless res.has_more
+        @done = true
+        close
+      end
+    end
+
+    def parse_row(raw)
+      row = {}
+      fields = raw.split(metadata.delim)
+
+      row_convertor.each do |c, p, i|
+        v = fields[i]
+        row[c] = (p ? p.call(v) : v unless v == NULL)
+      end
+
+      row
+    end
+
+    def row_convertor
+      @row_convertor ||= columns.zip(metadata.schema.fieldSchemas.map{|s| typecast_map[s.type]}, (0...(columns.length)).to_a)
+    end
+  end
+end

data/lib/impala/protocol.rb

@@ -0,0 +1,8 @@
+require 'impala/protocol/impala_service'
+
+module Impala
+  # Taken as a whole, this module contains all the thrift-generated stuff that
+  # defines the Impala protocol.
+  module Protocol
+  end
+end

data/lib/impala/protocol/beeswax_constants.rb

@@ -0,0 +1,15 @@
+#
+# Autogenerated by Thrift Compiler (0.9.1)
+#
+# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING
+#
+
+require 'thrift'
+require 'beeswax_types'
+
+module Impala
+  module Protocol
+    module Beeswax
+    end
+  end
+end