<commit_before># suppress warnings
old_verbose, $VERBOSE = $VERBOSE, nil

raise 'Thrift is not loaded' unless defined?(Thrift)
raise 'RBHive is not loaded' unless defined?(RBHive)

# require thrift autogenerated files
require File.join(File.dirname(__FILE__), *%w[.. thrift t_c_l_i_service_constants])
require File.join(File.dirname(__FILE__), *%w[.. thrift t_c_l_i_service])
require File.join(File.dirname(__FILE__), *%w[.. thrift sasl_client_transport])

# restore warnings
$VERBOSE = old_verbose


module RBHive
  def tcli_connect(server, port=10_000, sasl_params={})
    connection = RBHive::TCLIConnection.new(server, port, sasl_params)
    ret = nil
    begin
      connection.open
      connection.open_session
      ret = yield(connection)
    ensure
      connection.close_session if connection.session?
      connection.close
      ret
    end
  end
  module_function :tcli_connect

  class StdOutLogger
    %w(fatal error warn info debug).each do |level|
      define_method level.to_sym do |message|
        STDOUT.puts(message)
     end
   end
  end

  class TCLIConnection
    attr_reader :client

    def initialize(server, port=10_000, sasl_params={}, logger=StdOutLogger.new)
      @socket = Thrift::Socket.new(server, port)
      @socket.timeout = 1800
      @logger = logger
      if sasl_params.present?
        @logger.info("Initializing transport with SASL support")
        @transport = Thrift::SaslClientTransport.new(@socket, sasl_params)
      else
        @transport = Thrift::BufferedTransport.new(@socket)
      end
      @protocol = Thrift::BinaryProtocol.new(@transport)
      @client = TCLIService::Client.new(@protocol)
      @session = nil
      @logger.info("Connecting to HiveServer2 #{server} on port #{port}")
      @mutex = Mutex.new
    end

    def open
      @transport.open
    end

    def close
      @transport.close
    end

    def open_session
      @session = @client.OpenSession(prepare_open_session)
    end

    def close_session
      @client.CloseSession prepare_close_session
      @session = nil
    end

    def session?
      @session && @session.sessionHandle
    end

    def client
      @client
    end

    def execute(query)
      execute_safe(query)
    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}")
      self.execute("SET #{name}=#{value}")
    end

    def fetch(query)
      safe do
        op_handle = execute_unsafe(query).operationHandle
        fetch_req = prepare_fetch_results(op_handle)
        fetch_results = client.FetchResults(fetch_req)
        raise fetch_results.status.try(:errorMessage, 'Execution failed!').to_s if fetch_results.status.statusCode != 0
        rows = fetch_results.results.rows
        the_schema = TCLISchemaDefinition.new(get_schema_for( op_handle ), rows.first)
        TCLIResultSet.new(rows, the_schema)
      end
    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}")
      req = prepare_execute_statement(query)
      client.ExecuteStatement(req)
    end

    def safe
      ret = nil
      @mutex.synchronize { ret = yield }
      ret
    end

    def prepare_open_session
      TOpenSessionReq.new
    end

    def prepare_close_session
      TCloseSessionReq.new( sessionHandle: @session.sessionHandle )
    end

    def prepare_execute_statement(query)
      TExecuteStatementReq.new( sessionHandle: @session.sessionHandle, statement: query.to_s )<commit_msg>minor cosmetic changes in openSession request and session<commit_after># suppress warnings
old_verbose, $VERBOSE = $VERBOSE, nil

raise 'Thrift is not loaded' unless defined?(Thrift)
raise 'RBHive is not loaded' unless defined?(RBHive)

# require thrift autogenerated files
require File.join(File.dirname(__FILE__), *%w[.. thrift t_c_l_i_service_constants])
require File.join(File.dirname(__FILE__), *%w[.. thrift t_c_l_i_service])
require File.join(File.dirname(__FILE__), *%w[.. thrift sasl_client_transport])

# restore warnings
$VERBOSE = old_verbose


module RBHive
  def tcli_connect(server, port=10_000, sasl_params={})
    connection = RBHive::TCLIConnection.new(server, port, sasl_params)
    ret = nil
    begin
      connection.open
      connection.open_session
      ret = yield(connection)
    ensure
      connection.close_session if connection.session
      connection.close
      ret
    end
  end
  module_function :tcli_connect

  class StdOutLogger
    %w(fatal error warn info debug).each do |level|
      define_method level.to_sym do |message|
        STDOUT.puts(message)
     end
   end
  end

  class TCLIConnection
    attr_reader :client

    def initialize(server, port=10_000, sasl_params=nil, logger=StdOutLogger.new)
      @socket = Thrift::Socket.new(server, port)
      @socket.timeout = 1800
      @logger = logger
      if sasl_params.present?
        @logger.info("Initializing transport with SASL support")
        @transport = Thrift::SaslClientTransport.new(@socket, sasl_params)
        @sasl_params = case sasl_params
                       when Hash then sasl_params.symbolize_keys
                       when Hashie::Mash then sasl_params.to_hash(symbolize_keys: true)
                       else nil
                       end
      else
        @transport = Thrift::BufferedTransport.new(@socket)
      end
      @protocol = Thrift::BinaryProtocol.new(@transport)
      @client = TCLIService::Client.new(@protocol)
      @session = nil
      @logger.info("Connecting to HiveServer2 #{server} on port #{port}")
      @mutex = Mutex.new
    end

    def open
      @transport.open
    end

    def close
      @transport.close
    end

    def open_session
      @session = @client.OpenSession(prepare_open_session)
    end

    def close_session
      @client.CloseSession prepare_close_session
      @session = nil
    end

    def session
      @session && @session.sessionHandle
    end

    def client
      @client
    end

    def execute(query)
      execute_safe(query)
    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}")
      self.execute("SET #{name}=#{value}")
    end

    def fetch(query)
      safe do
        op_handle = execute_unsafe(query).operationHandle
        fetch_req = prepare_fetch_results(op_handle)
        fetch_results = client.FetchResults(fetch_req)
        raise fetch_results.status.try(:errorMessage, 'Execution failed!').to_s if fetch_results.status.statusCode != 0
        rows = fetch_results.results.rows
        the_schema = TCLISchemaDefinition.new(get_schema_for( op_handle ), rows.first)
        TCLIResultSet.new(rows, the_schema)
      end
    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}")
      req = prepare_execute_statement(query)
      client.ExecuteStatement(req)
    end

    def safe
      ret = nil
      @mutex.synchronize { ret = yield }
      ret
    end

    def prepare_open_session
      TOpenSessionReq.new( @sasl_params.presence )
    end

    def prepare_close_session
      TCloseSessionReq.new( sessionHandle: self.session )
    end

    def prepare_execute_statement(query)
      TExecuteStatementReq.new( sessionHandle: self.session, statement: query.to_s )
{'fraction_non_alphanumeric': 0.09229390681003584, 'fraction_numerical': 0.00409626216077829, 'mean_word_length': 2.920220772704466, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2954118', 'n_tokens_mistral': 2460, 'n_tokens_neox': 2383, 'n_words': 541}