sequel_impala 1.1.0

data/lib/rbhive/t_c_l_i_result_set.rb

@@ -0,0 +1,3 @@
+module RBHive
+  class TCLIResultSet < ResultSet; end
+end

data/lib/rbhive/t_c_l_i_schema_definition.rb

@@ -0,0 +1,87 @@
+require 'json'
+
+module RBHive
+  class TCLISchemaDefinition
+    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,
+      :float    => :to_f,
+      :double   => :to_f,
+      :int      => :to_i,
+      :bigint   => :to_i,
+      :smallint => :to_i,
+      :tinyint  => :to_i,
+    }
+
+    def initialize(schema, example_row)
+      @schema = schema
+      @example_row = example_row ? example_row.colVals : []
+    end
+
+    def column_names
+      @column_names ||= begin
+        schema_names = @schema.columns.map {|c| c.columnName }
+
+        # 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.columns.find {|s| s.columnName.to_sym == c }
+        # If the column isn't in the schema (eg partitions in SELECT * queries) assume they are strings
+        type = TYPE_NAMES[definition.typeDesc.types.first.primitiveEntry.type].downcase rescue nil
+        hsh[c] = definition && type ? type.to_sym : :string
+        hsh
+      end
+    end
+
+    def coerce_row(row)
+      column_names.zip(row.colVals.map(&:get_value).map(&:value)).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/table_schema.rb

@@ -0,0 +1,122 @@
+module RBHive
+  class TableSchema
+    attr_accessor :name
+    attr_reader :columns, :partitions
+    def initialize(name, comment=nil, options={}, &blk)
+      @name, @comment = name, comment
+      @location = options[:location] || nil
+      @field_sep = options[:field_sep] || "\t"
+      @line_sep = options[:line_sep] || "\n"
+      @collection_sep = options[:collection_sep] || "|"
+      @stored_as = options[:stored_as] || :textfile
+      @columns = []
+      @partitions = []
+      @serde_name = nil
+      @serde_properties = {}
+      instance_eval(&blk) if blk
+    end
+
+    def column(name, type, comment=nil)
+      @columns << Column.new(name, type, comment)
+    end
+
+    def partition(name, type, comment=nil)
+      @partitions << Column.new(name, type, comment)
+    end
+
+    def serde(name, properties={})
+      @serde_name = name
+      @serde_properties = properties
+    end
+
+    def create_table_statement()
+      %[CREATE #{external}TABLE #{table_statement}
+  ROW FORMAT #{row_format_statement}
+  STORED AS #{stored_as}
+  #{location}]
+    end
+
+    def stored_as
+      @stored_as.to_s.upcase
+    end
+
+    def row_format_statement
+      if @serde_name
+        serde_statement
+      else
+        delimited_statement
+      end
+    end
+
+    def delimited_statement
+      %(DELIMITED
+  FIELDS TERMINATED BY '#{@field_sep}'
+  COLLECTION ITEMS TERMINATED BY '#{@collection_sep}'
+  LINES TERMINATED BY '#{@line_sep}')
+    end
+
+    def serde_statement
+      %(SERDE '#{@serde_name}'\n#{serde_properties_statement})
+    end
+
+    def serde_properties_statement
+      return '' unless @serde_properties.any?
+      kvs = @serde_properties.map { |k,v| %("#{k}" = "#{v}") }.join(",\n")
+      %(WITH SERDEPROPERTIES (#{kvs}))
+    end
+
+    def replace_columns_statement
+      alter_columns_statement("REPLACE")
+    end
+
+    def add_columns_statement
+      alter_columns_statement("ADD")
+    end
+
+    def to_s
+      table_statement
+    end
+
+    private
+
+    def external
+      @location.nil? ? '' : 'EXTERNAL '
+    end
+
+    def table_statement
+      comment_string = (@comment.nil? ? '' : " COMMENT '#{@comment}'")
+      %[`#{@name}` #{column_statement}#{comment_string}\n#{partition_statement}]
+    end
+
+    def location
+      @location.nil? ? '' : "LOCATION '#{@location}'"
+    end
+
+    def alter_columns_statement(add_or_replace)
+      %[ALTER TABLE `#{name}` #{add_or_replace} COLUMNS #{column_statement}]
+    end
+
+    def column_statement
+      cols = @columns.join(",\n")
+      "(\n#{cols}\n)"
+    end
+
+    def partition_statement
+      return "" if @partitions.nil? || @partitions.empty?
+      cols = @partitions.join(",\n")
+      "PARTITIONED BY (\n#{cols}\n)"
+    end
+
+    class Column
+      attr_reader :name, :type, :comment
+      def initialize(name, type, comment=nil)
+        @name, @type, @comment = name, type, comment
+      end
+
+      def to_s
+        comment_string = @comment.nil? ? '' : " COMMENT '#{@comment}'"
+        "`#{@name}` #{@type.to_s.upcase}#{comment_string}"
+      end
+    end
+  end
+end

data/lib/rbhive/version.rb

@@ -0,0 +1,3 @@
+module RBHive
+  VERSION = '1.0.3.pre'
+end

data/lib/sequel/adapters/impala.rb

@@ -0,0 +1,220 @@
+require 'impala'
+require 'sequel/adapters/shared/impala'
+
+module Sequel
+  module Impala
+    class Database < Sequel::Database
+      RECORD_QUERY_PROFILE = Object.new.freeze
+
+      include DatabaseMethods
+
+      # Exception classes used by Impala.
+      ImpalaExceptions = [
+        ::Impala::InvalidQueryError,
+        ::Impala::ConnectionError,
+        ::Impala::CursorError,
+        ::Impala::ParsingError,
+        ::Impala::Protocol::Beeswax::BeeswaxException,
+        ::Thrift::TransportException,
+        IOError
+      ].freeze
+
+      DisconnectExceptions = [
+        ::Thrift::TransportException,
+        IOError
+      ].freeze
+
+      set_adapter_scheme :impala
+
+      # Connect to the Impala server.  Currently, only the :host and :port options
+      # are respected, and they default to 'localhost' and 21000, respectively.
+      def connect(server)
+        opts = server_opts(server)
+        force_database(::Impala.connect(opts[:host]||'localhost', (opts[:port]||21000).to_i, opts), opts[:database])
+      end
+
+      def database_error_classes
+        ImpalaExceptions
+      end
+
+      def disconnect_connection(c)
+        log_info("Closing connection: #{c}")
+        c.close
+      rescue *DisconnectExceptions
+      end
+
+      def execute(sql, opts=OPTS)
+        synchronize(opts[:server]) do |c|
+          # here's my super-hack to get DDL calls to record their profiles and query_ids
+          opts = self.opts.select { |k, v| [:query_id_name, :profile_name].include?(k) }.merge(opts)
+          begin
+            cursor = record_query_id(opts) do
+              log_connection_yield(sql, c) do
+                c.execute(sql){}
+              end
+            end
+            yield cursor if block_given?
+            nil
+          rescue *ImpalaExceptions => e
+            raise_error(e)
+          ensure
+            record_profile(cursor, opts)
+            log_info("Closing cursor: #{cursor.inspect}")
+            log_query_url(cursor.handle) if cursor && cursor.handle
+            cursor.close if cursor && cursor.open?
+          end
+        end
+      end
+
+      def log_query_url(handle)
+        log_info(sprintf(ENV['SEQUEL_IMPALA_QUERY_URL'], query_id: handle.id)) if ENV['SEQUEL_IMPALA_QUERY_URL']
+      end
+
+      def query_id_and_profile(query_id_name=:default, profile_name=:default)
+        key = RECORD_QUERY_PROFILE
+        prev_profile_name = prev_query_id_name = nil
+        begin
+          Sequel.synchronize do
+            prev_query_id_name = @query_ids[key]
+            prev_profile_name = @runtime_profiles[key]
+            @query_ids[key] = query_id_name
+            @runtime_profiles[key] = profile_name
+          end
+
+          yield
+        ensure
+          Sequel.synchronize do
+            @query_ids[key] = prev_query_id_name
+            @runtime_profiles[key] = prev_profile_name
+          end
+        end
+      end
+
+      def profile_for(profile_name=:default)
+        Sequel.synchronize{@runtime_profiles[profile_name]}
+      end
+
+      def query_id_for(query_id_name=:default)
+        Sequel.synchronize{@query_ids[query_id_name]}
+      end
+
+      private
+
+      def dataset_class_default
+        Dataset
+      end
+
+      def record_profile(cursor, opts)
+        if cursor && (profile_name = opts[:profile_name] || Sequel.synchronize{@runtime_profiles[RECORD_QUERY_PROFILE]})
+          profile = cursor.runtime_profile
+          Sequel.synchronize{@runtime_profiles[profile_name] = profile}
+        end
+      end
+
+      def record_query_id(opts = OPTS)
+        query_id_name = opts[:query_id_name] || Sequel.synchronize{@query_ids[RECORD_QUERY_PROFILE]}
+        start = Time.now if query_id_name
+
+        cursor = yield
+
+        if cursor && query_id_name
+          h = { query_id: cursor.handle.id, start_time: start }
+          Sequel.synchronize{ @query_ids[query_id_name] = h }
+        end
+
+        cursor
+      end
+
+      def adapter_initialize
+        @runtime_profiles = {}
+        @query_ids = {}
+      end
+
+      def connection_execute_method
+        :query
+      end
+
+      # Impala raises IOError if it detects a problem on the connection, and
+      # in most cases that results in an unusable connection, so treat it as a
+      # disconnect error so Sequel will reconnect.
+      def disconnect_error?(exception, opts)
+        case exception
+        when *DisconnectExceptions
+          true
+        else
+          super
+        end
+      end
+
+      # Use DESCRIBE to get the column names and types for the table.
+      def schema_parse_table(table_name, opts)
+        m = output_identifier_meth(opts[:dataset])
+
+        table = if opts[:schema]
+          Sequel.qualify(opts[:schema], table_name)
+        else
+          Sequel.identifier(table_name)
+        end
+
+        describe(table, opts).map do |row|
+          row[:db_type] = row[:type]
+          row[:type] = schema_column_type(row[:db_type])
+          row[:default] = nil
+          row[:primary_key] = false
+          [m.call(row.delete(:name)), row]
+        end
+      end
+    end
+
+    class Dataset < Sequel::Dataset
+      include DatasetMethods
+
+      APOS = "'".freeze
+      STRING_ESCAPES = {
+        "\\" => "\\\\".freeze,
+        "'" => "\\'".freeze,
+        "\n" => "\\n".freeze,
+        "\r" => "\\r".freeze,
+        "\0" => "\\0".freeze,
+        "\b" => "\\b".freeze,
+        "\04" => "\\Z".freeze,
+       # Impala is supposed to support this, but using it
+       # breaks things to the point of returning bad data.
+       # If you don't do this, the tabs in the input
+       # get converted to spaces, but that's better than the
+       # alternative.
+       # "\t" => "\\t".freeze,
+      }.freeze
+      STRING_ESCAPE_RE = /(#{Regexp.union(STRING_ESCAPES.keys)})/
+
+      def fetch_rows(sql)
+        execute(sql, @opts) do |cursor|
+          self.columns = cursor.columns.map!{|c| output_identifier(c)}
+          cursor.typecast_map['timestamp'] = db.method(:to_application_timestamp)
+          cursor.each do |row|
+            yield row
+          end
+        end
+
+        self
+      end
+
+      def profile(profile_name=:default)
+        clone(:profile_name => profile_name)
+      end
+
+      def query_id(query_id_name=:default)
+        clone(:query_id_name => query_id_name)
+      end
+
+      private
+
+      # Unlike the jdbc/hive2 driver, the impala driver requires you escape
+      # some values in string literals to get correct results, but not the
+      # tab character or things break.
+      def literal_string_append(sql, s)
+        sql << APOS << s.to_s.gsub(STRING_ESCAPE_RE){|m| STRING_ESCAPES[m]} << APOS
+      end
+    end
+  end
+end

data/lib/sequel/adapters/jdbc/hive2.rb

@@ -0,0 +1,36 @@
+require 'sequel/adapters/shared/impala'
+
+Sequel::JDBC.load_driver('org.apache.hive.jdbc.HiveDriver', :Hive2)
+
+module Sequel
+  module JDBC
+    Sequel.synchronize do
+      DATABASE_SETUP[:hive2] = proc do |db|
+        db.extend(Sequel::JDBC::Hive2::DatabaseMethods)
+        db.dataset_class = Sequel::JDBC::Hive2::Dataset
+        org.apache.hive.jdbc.HiveDriver
+      end
+    end
+
+    module Hive2
+      module DatabaseMethods
+        include Sequel::Impala::DatabaseMethods
+
+        # Recognize wrapped and unwrapped java.net.SocketExceptions as disconnect errors
+        def disconnect_error?(exception, opts)
+          super || exception.message =~ /\A(Java::JavaSql::SQLException: )?org\.apache\.thrift\.transport\.TTransportException: java\.net\.SocketException/
+        end
+
+        def disconnect_connection(c)
+          super
+        rescue java.sql.SQLException
+          nil
+        end
+      end
+
+      class Dataset < JDBC::Dataset
+        include Sequel::Impala::DatasetMethods
+      end
+    end
+  end
+end