sequel_impala 1.1.0

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/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,55 @@
+
+# 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'
+begin
+require 'impala/sasl_transport'
+rescue LoadError
+  # gssapi not supported by operating system, continue as impala adapter
+  # can be used in buffered (non-SASL) mode.
+end
+require 'impala/progress_reporter'
+require 'impala/cursor'
+require 'impala/connection'
+require 'impala/thrift_patch'
+
+module Impala
+  DEFAULT_HOST = 'localhost'
+  DEFAULT_PORT = 21000
+  class InvalidQueryError < StandardError; end
+  class ConnectionError < StandardError; end
+  class CursorError < StandardError; end
+  class ParsingError < StandardError; 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, options={})
+    connection = Connection.new(host, port, options)
+
+    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,180 @@
+module Impala
+  # This object represents a connection to an Impala server. It can be used to
+  # perform queries on the database.
+  class Connection
+    attr_accessor :host, :port
+
+    # Don't instantiate Connections directly; instead, use {Impala.connect}.
+    def initialize(host, port, options={})
+      @host = host
+      @port = port
+      @connected = false
+      @options = options.dup
+      @options[:transport] ||= :buffered
+      @loggers = @options.fetch(:loggers, [])
+      open
+    end
+
+    def inspect
+      "#<#{self.class} #{@host}:#{@port}#{open? ? '' : ' (DISCONNECTED)'}>"
+    end
+
+    # Open the connection if it's currently closed.
+    def open
+      return if @connected
+
+      @transport = thrift_transport(host, port)
+      @transport.open do |transport|
+        enable_keepalive(transport)
+      end
+
+      proto = Thrift::BinaryProtocol.new(@transport)
+      @service = Protocol::ImpalaService::Client.new(proto)
+      @connected = true
+    end
+
+    def thrift_transport(server, port)
+      socket = thrift_socket(server, port, @options[:timeout])
+
+      case @options[:transport]
+      when :buffered
+        return Thrift::BufferedTransport.new(socket)
+      when :sasl
+        opts = parse_sasl_params(@options[:sasl_params])
+        mechanism = opts.delete(:mechanism)
+        return SASLTransport.new(socket, mechanism, opts)
+      else
+        raise "Unrecognised transport type '#{@options[: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
+
+    # 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)
+
+      cursor = Cursor.new(handle, @service, @options)
+      cursor.wait!
+      cursor
+    end
+
+    def close_handle(handle)
+      @service.close(handle)
+    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.query(query)
+    end
+
+    def enable_keepalive(transport)
+      s = transport.handle
+      log_debug("Enabling KEEPALIVE...")
+      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)
+        opts = {}
+
+        if defined?(::Socket::TCP_KEEPIDLE)
+          opts[::Socket::TCP_KEEPIDLE] = 60
+        end
+
+        if defined?(::Socket::TCP_KEEPINTVL)
+          opts[::Socket::TCP_KEEPINTVL] = 10
+        end
+
+        if defined?(::Socket::TCP_KEEPCNT)
+          opts[::Socket::TCP_KEEPCNT] = 5
+        end
+
+        log_debug("Also enabling: #{opts.inspect}")
+        opts.each do |opt, value|
+          s.setsockopt(::Socket::SOL_TCP, opt, value)
+        end
+      end
+    end
+
+    def log_debug(message)
+      @loggers.each do |logger|
+        logger.debug(message)
+      end
+    end
+  end
+end

data/lib/impala/cursor.rb

@@ -0,0 +1,200 @@
+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 :typecast_map
+
+    attr_reader :handle
+
+    def initialize(handle, service, options = {})
+      @handle = handle
+      @service = service
+
+      @row_buffer = []
+      @done = false
+      @open = true
+      @typecast_map = TYPECAST_MAP.dup
+      @options = options.dup
+      @progress_reporter = ProgressReporter.new(self, @options)
+      @poll_every = options.fetch(:poll_every, 0.1)
+    end
+
+    def columns
+      @columns ||= metadata.schema.fieldSchemas.map(&:name)
+    end
+
+    def inspect
+      "#<#{self.class}#{handle ? " QueryID: #{handle.id}" : ''}#{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 the query is done running, and results can be fetched.
+    def query_done?
+      [
+        Protocol::Beeswax::QueryState::EXCEPTION,
+        Protocol::Beeswax::QueryState::FINISHED
+      ].include?(@service.get_state(@handle))
+    end
+
+    # Blocks until the query done running.
+    def wait!
+      until query_done?
+        periodic_callback
+        sleep @poll_every
+      end
+    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
+
+    def exec_summary
+      @service.GetExecSummary(@handle)
+    end
+
+    # Returns the progress for the query.
+    def progress
+      summary = exec_summary
+      summary.progress.num_completed_scan_ranges.to_f / summary.progress.total_scan_ranges.to_f
+    end
+
+    private
+
+    attr :progress_reporter
+
+    def periodic_callback
+      return unless progress_reporter.show?
+      progress_reporter.report
+    end
+
+    def metadata
+      @metadata ||= @service.get_results_metadata(@handle)
+    end
+
+    def fetch_more
+      fetch_batch until @done || @row_buffer.count >= BUFFER_SIZE
+    end
+
+    def exceptional?
+      @service.get_state(@handle) == Protocol::Beeswax::QueryState::EXCEPTION
+    end
+
+    def fetch_batch
+      raise CursorError.new("Cursor has expired or been closed") unless @open
+      raise ConnectionError.new("The query was aborted") if exceptional?
+
+      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