sequel-impala 1.0.1 → 1.1.0

data/lib/rbhive/t_c_l_i_schema_definition.rb

@@ -48,7 +48,7 @@ module RBHive
     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.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

data/lib/rbhive/table_schema.rb

@@ -119,4 +119,4 @@ module RBHive
       end
     end
   end
-end
+end

data/lib/sequel/adapters/impala.rb

@@ -8,7 +8,10 @@ module Sequel
 
       # Exception classes used by Impala.
       ImpalaExceptions = [
-        ::Impala::Error,
+        ::Impala::InvalidQueryError,
+        ::Impala::ConnectionError,
+        ::Impala::CursorError,
+        ::Impala::ParsingError,
         ::Impala::Protocol::Beeswax::BeeswaxException,
         ::Thrift::TransportException,
         IOError
@@ -21,11 +24,12 @@ module Sequel
 
       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)
-        ::Impala.connect(opts[:host]||'localhost', (opts[:port]||21000).to_i)
+        ::Impala.connect(opts[:host]||'localhost', (opts[:port]||21000).to_i, opts)
       end
 
       def database_error_classes
@@ -40,19 +44,62 @@ module Sequel
       def execute(sql, opts=OPTS)
         synchronize(opts[:server]) do |c|
           begin
-            cursor = log_yield(sql){c.execute(sql)}
+            cursor = record_query_id(opts) do
+              log_yield(sql) do
+                c.execute(sql){}
+              end
+            end
             yield cursor if block_given?
             nil
           rescue *ImpalaExceptions => e
+            puts $!.message
+            puts $!.backtrace.join("\n")
             raise_error(e)
+          rescue
+            puts $!.message
+            puts $!.backtrace.join("\n")
+            raise
           ensure
+            record_profile(cursor, opts)
             cursor.close if cursor && cursor.open?
           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 record_profile(cursor, opts)
+        if cursor && profile_name = opts[:profile_name]
+          profile = cursor.runtime_profile
+          Sequel.synchronize{@runtime_profiles[profile_name] = profile}
+        end
+      end
+
+      def record_query_id(opts = OPTS)
+        start = Time.now if opts[:query_id_name]
+
+        cursor = yield
+
+        if cursor && query_id_name = opts[:query_id_name]
+          Sequel.synchronize{ @query_ids[query_id_name] = { query_id: cursor.handle.id, start_time: start } }
+        end
+
+        cursor
+      end
+
+      def adapter_initialize
+        @runtime_profiles = {}
+        @query_ids = {}
+      end
+
       def connection_execute_method
         :query
       end
@@ -113,13 +160,23 @@ module Sequel
       STRING_ESCAPE_RE = /(#{Regexp.union(STRING_ESCAPES.keys)})/
 
       def fetch_rows(sql)
-        execute(sql) do |cursor|
+        execute(sql, @opts) do |cursor|
           @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
@@ -128,8 +185,8 @@ module Sequel
       # 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 
+        sql << APOS << s.to_s.gsub(STRING_ESCAPE_RE){|m| STRING_ESCAPES[m]} << APOS
       end
     end
   end
-end 
+end

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

@@ -17,7 +17,7 @@ module Sequel
         extend Sequel::Database::ResetIdentifierMangling
         include Sequel::Impala::DatabaseMethods
       end
-      
+
       class Dataset < JDBC::Dataset
         include Sequel::Impala::DatasetMethods
       end

data/lib/sequel/adapters/rbhive.rb

@@ -74,6 +74,7 @@ module Sequel
       def execute(sql, opts=OPTS)
         synchronize(opts[:server]) do |c|
           begin
+            puts sql
             r = log_yield(sql){c.execute(sql)}
             yield(c, r) if block_given?
             nil
@@ -166,9 +167,9 @@ module Sequel
       private
 
       def literal_string_append(sql, s)
-        sql << APOS << s.to_s.gsub(STRING_ESCAPE_RE){|m| STRING_ESCAPES[m]} << APOS 
+        sql << APOS << s.to_s.gsub(STRING_ESCAPE_RE){|m| STRING_ESCAPES[m]} << APOS
       end
     end
   end
-end 
+end
 

data/lib/sequel/adapters/shared/impala.rb

@@ -1,6 +1,10 @@
 module Sequel
   module Impala
+    Sequel::Database.set_shared_adapter_scheme :impala, self
+
     module DatabaseMethods
+      extend Sequel::Database::ResetIdentifierMangling
+
       # Do not use a composite primary key, foreign keys, or an
       # index when creating a join table, as Impala doesn't support those.
       def create_join_table(hash, options=OPTS)
@@ -12,6 +16,14 @@ module Sequel
         end
       end
 
+      def refresh(table_name)
+        run(refresh_sql(table_name))
+      end
+
+      def compute_stats(table_name)
+        run(compute_stats_sql(table_name))
+      end
+
       # Create a database/schema in Imapala.
       #
       # Options:
@@ -98,7 +110,12 @@ module Sequel
           if schema = search_path_table_schemas[table]
             Sequel.qualify(schema, table)
           else
-            Sequel.identifier(table)
+            invalidate_table_schemas
+            if schema = search_path_table_schemas[table]
+              Sequel.qualify(schema, table)
+            else
+              Sequel.identifier(table)
+            end
           end
         when SQL::Identifier
           implicit_qualify(table.value.to_s)
@@ -150,7 +167,7 @@ module Sequel
       # tables and not views.  The Hive2 JDBC driver returns views when listing
       # tables and nothing when listing views.
       def tables(opts=OPTS)
-        _tables(opts).select{|t| is_valid_table?(t)}
+        _tables(opts).select{|t| is_valid_table?(t, opts)}
       end
 
       # Impala doesn't support transactions, so instead of issuing a
@@ -166,7 +183,30 @@ module Sequel
       # Determine the available views for listing all tables via JDBC (which
       # includes both tables and views), and removing all valid tables.
       def views(opts=OPTS)
-        _tables(opts).reject{|t| is_valid_table?(t)}
+        _tables(opts).reject{|t| is_valid_table?(t, opts)}
+      end
+
+      def invalidate_table_schemas
+        @search_path_table_schemas = nil
+      end
+
+      # Creates a dataset that uses the VALUES clause:
+      #
+      #   DB.values([[1, 2], [3, 4]])
+      #   VALUES ((1, 2), (3, 4))
+      def values(v)
+        @default_dataset.clone(:values=>v)
+      end
+
+      def invalidate_table_schemas
+        @search_path_table_schemas = nil
+      end
+
+      # Sets options in the current db connection for each key/value pair
+      def set(opts)
+        set_sql(opts).each do |sql|
+          run(sql)
+        end
       end
 
       private
@@ -251,6 +291,14 @@ module Sequel
         sql
       end
 
+      def refresh_sql(table_name)
+        "REFRESH #{quote_schema_table(table_name)}"
+      end
+
+      def compute_stats_sql(table_name)
+        "COMPUTE STATS #{quote_schema_table(table_name)}"
+      end
+
       def drop_schema_sql(schema, options)
         "DROP SCHEMA #{'IF EXISTS ' if options[:if_exists]}#{quote_identifier(schema)}"
       end
@@ -264,6 +312,10 @@ module Sequel
         nil
       end
 
+      def quote_identifiers_default
+        true
+      end
+
       def search_path_table_schemas
         @search_path_table_schemas ||= begin
           search_path = opts[:search_path]
@@ -280,7 +332,8 @@ module Sequel
 
       # SHOW TABLE STATS will raise an error if given a view and not a table,
       # so use that to differentiate tables from views.
-      def is_valid_table?(t)
+      def is_valid_table?(t, opts=OPTS)
+        t = [opts[:schema], t].map(&:to_s).join('__').to_sym if opts[:schema]
         rows = describe(t, :formatted=>true)
         if row = rows.find{|r| r[:name].to_s.strip == 'Table Type:'}
           row[:type].to_s.strip !~ /VIEW/
@@ -302,6 +355,21 @@ module Sequel
         )
       end
 
+      # Impala doesn't like the word "integer"
+      def type_literal_generic_integer(column)
+        :int
+      end
+
+      # Impala doesn't like the word "biginteger"
+      def type_literal_generic_bignum_symbol(column)
+        :bigint
+      end
+
+      # Impala doesn't like the word "biginteger"
+      def type_literal_generic_bignum(column)
+        :bigint
+      end
+
       # Impala doesn't support date columns yet, so use timestamp until date
       # is natively supported.
       def type_literal_generic_date(column)
@@ -331,6 +399,10 @@ module Sequel
           :string
         end
       end
+
+      def set_sql(opts)
+        opts.map { |k, v| "SET #{k}=#{v}" }
+      end
     end
 
     module DatasetMethods
@@ -347,8 +419,10 @@ module Sequel
       NOT = 'NOT '.freeze
       REGEXP = ' REGEXP '.freeze
       EXCEPT_SOURCE_COLUMN = :__source__
+      EXCEPT_STRATEGIES = [:not_exists, :not_in, :left_join, :group_by].freeze
+      SELECT_VALUES = 'VALUES '.freeze
 
-      Dataset.def_sql_method(self, :select, %w'with select distinct columns from join where group having compounds order limit')
+      Dataset.def_sql_method(self, :select, [['if opts[:values]', %w'values'], ['else', %w'with select distinct columns from join where group having compounds order limit']])
 
       # Handle string concatenation using the concat string function.
       # Don't use the ESCAPE syntax when using LIKE/NOT LIKE, as
@@ -454,19 +528,48 @@ module Sequel
         get(Sequel::SQL::AliasedExpression.new(1, :one)).nil?
       end
 
-      # Emulate INTERSECT using a UNION ALL and checking for values in only the first table.
+      # Emulate EXCEPT using a chosen strategy and checking for values in only the first table.
       def except(other, opts=OPTS)
         raise(InvalidOperation, "EXCEPT ALL not supported") if opts[:all]
         raise(InvalidOperation, "The :from_self=>false option to except is not supported") if opts[:from_self] == false
-        cols = columns
-        rhs = other.from_self.select_group(*other.columns).select_append(Sequel.expr(2).as(EXCEPT_SOURCE_COLUMN))
-        from_self.
-          select_group(*cols).
-          select_append(Sequel.expr(1).as(EXCEPT_SOURCE_COLUMN)).
-          union(rhs, all: true).
-          select_group(*cols).
-          having{{count{}.* => 1, min(EXCEPT_SOURCE_COLUMN) => 1}}.
-          from_self(opts)
+
+        strategy, *keys = @opts[:except_strategy]
+        ds = from_self(:alias=>:t1)
+
+        ds = case strategy
+        when :not_exists
+          ds.exclude(other.
+              from_self(:alias=>:t2).
+              where(keys.map{|key| [Sequel.qualify(:t1, key), Sequel.qualify(:t2, key)]}).
+              select(nil).
+              exists)
+        when :not_in
+          raise Sequel::Error, ":not_in EXCEPT strategy only supports a single key" unless keys.length == 1
+          key = keys.first
+          ds.exclude(Sequel.qualify(:t1, key)=>other.from_self(:alias=>:t2).select(key))
+        when :left_join
+          ds.left_join(other.from_self(:alias=>:t2).as(:t2), keys.map{|key| [key, key]}).
+            where(Sequel.or(keys.map{|key| [Sequel.qualify(:t2, key), nil]})).
+            select_all(:t1)
+        else
+          cols = columns
+          rhs = other.from_self.select_group(*other.columns).select_append(Sequel.expr(2).as(EXCEPT_SOURCE_COLUMN))
+          ds.select_group(*cols).
+            select_append(Sequel.expr(1).as(EXCEPT_SOURCE_COLUMN)).
+            union(rhs, all: true).
+            select_group(*cols).
+            having{{count{}.* => 1, min(EXCEPT_SOURCE_COLUMN) => 1}}
+        end
+
+        ds.from_self(opts)
+      end
+
+      # The strategy to use for EXCEPT emulation. By default, uses a GROUP BY emulation,
+      # as that doesn't require you provide a key column, but you can use this to choose
+      # a NOT EXISTS, NOT IN, or LEFT JOIN emulation, providing the unique key column.
+      def except_strategy(strategy, *keys)
+        raise Sequel::Error, "invalid EXCEPT strategy: #{strategy.inspect}" unless EXCEPT_STRATEGIES.include?(strategy)
+        clone(:except_strategy=>[strategy, *keys])
       end
 
       # Use INSERT OVERWRITE instead of INSERT INTO when inserting into this dataset:
@@ -487,19 +590,17 @@ module Sequel
         raise(InvalidOperation, "INTERSECT ALL not supported") if opts[:all]
         raise(InvalidOperation, "The :from_self=>false option to intersect is not supported") if opts[:from_self] == false
         cols = columns
-        from_self.
-          select_group(*cols).
-          union(other.from_self.select_group(*other.columns), all: true).
-          select_group(*cols).
-          having{count{}.* > 1}.
-          from_self(opts)
+        (from_self(alias: :l)
+          .join(other, cols)
+          .select_all(:l))
+          .from_self(opts)
       end
 
       # Impala supports non-recursive common table expressions.
       def supports_cte?(type=:select)
         true
       end
-      
+
       def supports_cte_in_subqueries?
         true
       end
@@ -520,7 +621,7 @@ module Sequel
       def supports_is_true?
         false
       end
-    
+
       # Impala doesn't support IN when used with multiple columns.
       def supports_multiple_column_in?
         false
@@ -602,7 +703,7 @@ module Sequel
       def insert_empty_columns_values
         [[columns.last], [nil]]
       end
-    
+
       def literal_true
         BOOL_TRUE
       end
@@ -619,7 +720,7 @@ module Sequel
       # Double backslashes in all strings, and escape all apostrophes with
       # backslashes.
       def literal_string_append(sql, s)
-        sql << APOS << s.to_s.gsub(STRING_ESCAPE_RE, STRING_ESCAPE_REPLACE) << APOS 
+        sql << APOS << s.to_s.gsub(STRING_ESCAPE_RE, STRING_ESCAPE_REPLACE) << APOS
       end
 
       def multi_insert_sql_strategy
@@ -638,6 +739,13 @@ module Sequel
         return unless opts[:from]
         super
       end
+
+
+      # Support VALUES clause instead of the SELECT clause to return rows.
+      def select_values_sql(sql)
+        sql << SELECT_VALUES
+        expression_list_append(sql, opts[:values])
+      end
     end
   end
 end

data/lib/thrift/sasl_client_transport.rb

@@ -21,7 +21,7 @@ module Thrift
       @sasl_mechanism = sasl_params.fetch(:mechanism, 'PLAIN')
       raise 'Unknown SASL mechanism: #{@sasl_mechanism}' unless ['PLAIN', 'GSSAPI'].include? @sasl_mechanism
       if @sasl_mechanism == 'GSSAPI'
-	require 'gssapi'
+        require 'gssapi'
         @sasl_remote_principal = sasl_params[:remote_principal]
         @sasl_remote_host = sasl_params[:remote_host]
         @gsscli = GSSAPI::Simple.new(@sasl_remote_host, @sasl_remote_principal)
@@ -34,7 +34,7 @@ module Thrift
       @index += sz
       ret = @rbuf.slice(@index - sz, sz) || Bytes.empty_byte_buffer
       if ret.length < sz
-	sz -= ret.length
+        sz -= ret.length
         read_into_buffer(@rbuf, [sz, len || 0].max)
         @index = sz
         ret += @rbuf.slice(0, sz) || Bytes.empty_byte_buffer