mirror of
https://github.com/apache/impala.git
synced 2025-12-22 11:28:09 -05:00
311 lines
11 KiB
Python
311 lines
11 KiB
Python
#!/usr/bin/env python
|
|
# Copyright (c) 2012 Cloudera, Inc. All rights reserved.
|
|
#
|
|
# Impala's shell
|
|
import cmd
|
|
import time
|
|
import sys
|
|
from optparse import OptionParser
|
|
|
|
from ImpalaService import ImpalaService
|
|
from JavaConstants.constants import DEFAULT_QUERY_OPTIONS
|
|
from ImpalaService.ImpalaService import TImpalaQueryOptions
|
|
from beeswaxd import BeeswaxService
|
|
from beeswaxd.BeeswaxService import QueryState
|
|
from thrift.transport.TSocket import TSocket
|
|
from thrift.transport.TTransport import TBufferedTransport, TTransportException
|
|
from thrift.protocol import TBinaryProtocol
|
|
from thrift.Thrift import TApplicationException
|
|
|
|
VERSION_STRING = "Impala v0.1 "
|
|
|
|
# Tarball / packaging build makes impala_build_version available
|
|
try:
|
|
from impala_build_version import get_version_string
|
|
from impala_build_version import get_build_date
|
|
VERSION_STRING += "(" + get_version_string()[:7] + ") built on " + get_build_date()
|
|
except Exception:
|
|
VERSION_STRING += "(build version not available)"
|
|
|
|
# Simple Impala shell. Can issue queries (with configurable options)
|
|
# Basic usage: type connect <host:port> to connect to an impalad
|
|
# Then issue queries or other commands. Tab-completion should show the set of
|
|
# available commands.
|
|
# TODO: (amongst others)
|
|
# - Column headers / metadata support
|
|
# - Report profiles
|
|
# - A lot of rpcs return a verbose TStatus from thrift/Status.thrift
|
|
# This will be useful for better error handling. The next iteration
|
|
# of the shell should handle this return paramter.
|
|
# - Make commands case agnostic.
|
|
class ImpalaShell(cmd.Cmd):
|
|
def __init__(self, options):
|
|
cmd.Cmd.__init__(self)
|
|
self.is_alive = True
|
|
self.disconnected_prompt = "[Not connected] > "
|
|
self.prompt = self.disconnected_prompt
|
|
self.connected = False
|
|
self.handle = None
|
|
self.impalad = None
|
|
self.imp_service = None
|
|
self.transport = None
|
|
self.query_options = {}
|
|
self.query_state = QueryState._NAMES_TO_VALUES
|
|
if options.impalad != None:
|
|
if self.do_connect(options.impalad) == False:
|
|
sys.exit(1)
|
|
|
|
def __make_default_options(self):
|
|
def get_name(option): return TImpalaQueryOptions._VALUES_TO_NAMES[option]
|
|
for option, default in DEFAULT_QUERY_OPTIONS.iteritems():
|
|
self.query_options[get_name(option)] = default
|
|
|
|
def __print_options(self):
|
|
print '\n'.join(["\t%s: %s" % (k,v) for (k,v) in self.query_options.iteritems()])
|
|
|
|
def __options_to_string_list(self):
|
|
return ["%s:%s" % (k,v) for (k,v) in self.query_options.iteritems()]
|
|
|
|
def do_options(self, args):
|
|
self.__print_options()
|
|
|
|
def do_set(self, args):
|
|
tokens = args.split(" ")
|
|
if len(tokens) != 2:
|
|
print "Error: SET <option> <value>"
|
|
return False
|
|
option_upper = tokens[0].upper()
|
|
if option_upper not in ImpalaService.TImpalaQueryOptions._NAMES_TO_VALUES.keys():
|
|
print "Unknown query option: %s" % (tokens[0],)
|
|
available_options = \
|
|
'\n\t'.join(ImpalaService.TImpalaQueryOptions._NAMES_TO_VALUES.keys())
|
|
print "Available query options are: \n\t%s" % available_options
|
|
return False
|
|
self.query_options[option_upper] = tokens[1]
|
|
return False
|
|
|
|
def do_quit(self, args):
|
|
print "Goodbye"
|
|
self.is_alive = False
|
|
return True
|
|
|
|
def do_connect(self, args):
|
|
tokens = args.split(" ")
|
|
if len(tokens) != 1:
|
|
print "CONNECT takes exactly one argument: <hostname:port> of impalad to connect to"
|
|
return False
|
|
try:
|
|
host, port = tokens[0].split(":")
|
|
self.impalad = (host, port)
|
|
except ValueError:
|
|
print "Connect string must be of form <hostname:port>"
|
|
return False
|
|
|
|
if self.__connect():
|
|
self.prompt = "[%s:%s] > " % self.impalad
|
|
|
|
def __connect(self):
|
|
if self.transport is not None:
|
|
self.transport.close()
|
|
self.transport = None
|
|
|
|
self.transport = TBufferedTransport(TSocket(self.impalad[0], int(self.impalad[1])))
|
|
try:
|
|
self.transport.open()
|
|
protocol = TBinaryProtocol.TBinaryProtocol(self.transport)
|
|
self.imp_service = ImpalaService.Client(protocol)
|
|
self.connected = True
|
|
except TTransportException, e:
|
|
print "Error connecting: %s" % (e,)
|
|
self.connected = False
|
|
return self.connected
|
|
|
|
def __query_with_results(self, query):
|
|
print "Query: %s" % (query.query,)
|
|
start, end = time.time(), 0
|
|
(self.handle, ok) = self.__do_rpc(lambda: self.imp_service.query(query))
|
|
fetching = False
|
|
if not ok: return False
|
|
|
|
try:
|
|
while self.__get_query_state() not in [self.query_state["FINISHED"],
|
|
self.query_state["EXCEPTION"]]:
|
|
time.sleep(0.05)
|
|
|
|
# The query finished with an exception, skip fetching results and close the handle.
|
|
if self.__get_query_state() == self.query_state["EXCEPTION"]:
|
|
print 'Query aborted, unable to fetch data'
|
|
return self.__close_query_handle()
|
|
# Results are ready, fetch them till they're done.
|
|
print 'Query finished, fetching results ...'
|
|
result_rows = []
|
|
fetching = True
|
|
while True:
|
|
# TODO: Fetch more than one row at a time.
|
|
(results, ok) = self.__do_rpc(lambda: self.imp_service.fetch(self.handle,
|
|
False, -1))
|
|
if not ok: return False
|
|
result_rows.extend(results.data)
|
|
if not results.has_more:
|
|
fetching = False
|
|
break
|
|
end = time.time()
|
|
except KeyboardInterrupt:
|
|
# If Ctrl^C was caught during query execution, cancel the query.
|
|
# If it was caught during fetching results, return to the prompt.
|
|
if not fetching:
|
|
return self.__cancel_query()
|
|
else:
|
|
print 'Fetching results aborted.'
|
|
return self.__close_query_handle()
|
|
|
|
print '\n'.join(result_rows)
|
|
print "Returned %d row(s) in %2.2fs" % (len(result_rows), end - start)
|
|
return self.__close_query_handle()
|
|
|
|
def __close_query_handle(self):
|
|
"""Close the query handle and get back to the prompt"""
|
|
self.__do_rpc(lambda: self.imp_service.close(self.handle))
|
|
self.handle = None
|
|
return False
|
|
|
|
def do_select(self, args):
|
|
"""Executes a SELECT... query, fetching all rows"""
|
|
query = BeeswaxService.Query()
|
|
query.query = "select %s" % (args,)
|
|
query.configuration = self.__options_to_string_list()
|
|
return self.__query_with_results(query)
|
|
|
|
def do_use(self, args):
|
|
"""Executes a USE... query"""
|
|
query = BeeswaxService.Query()
|
|
query.query = "use %s" % (args,)
|
|
query.configuration = self.__options_to_string_list()
|
|
return self.__query_with_results(query)
|
|
|
|
def do_show(self, args):
|
|
"""Executes a SHOW... query, fetching all rows"""
|
|
query = BeeswaxService.Query()
|
|
query.query = "show %s" % (args,)
|
|
query.configuration = self.__options_to_string_list()
|
|
return self.__query_with_results(query)
|
|
|
|
def do_describe(self, args):
|
|
"""Executes a DESCRIBE... query, fetching all rows"""
|
|
query = BeeswaxService.Query()
|
|
query.query = "describe %s" % (args,)
|
|
query.configuration = self.__options_to_string_list()
|
|
return self.__query_with_results(query)
|
|
|
|
def do_insert(self, args):
|
|
"""Executes an INSERT query"""
|
|
query = BeeswaxService.Query()
|
|
query.query = "insert %s" % (args,)
|
|
query.configuration = self.__options_to_string_list()
|
|
print "Query: %s" % (query.query,)
|
|
start, end = time.time(), 0
|
|
(handle, ok) = \
|
|
self.__do_rpc(lambda: self.imp_service.executeAndWait(query, "ImpalaCLI") )
|
|
if not ok: return False
|
|
(insert_result, ok) = self.__do_rpc(lambda: self.imp_service.CloseInsert(handle))
|
|
end = time.time()
|
|
if not ok: return False
|
|
num_rows = sum([int(k) for k in insert_result.rows_appended.values()])
|
|
print "Inserted %d rows in %2.2fs" % (num_rows, end - start)
|
|
|
|
def __cancel_query(self):
|
|
"""Cancel a query on keyboard interrupt from the shell."""
|
|
print 'Cancelling query ...'
|
|
# Cancel sets query_state to EXCEPTION before calling cancel() in the
|
|
# co-ordinator, so we don't need to wait.
|
|
self.__do_rpc(lambda: self.imp_service.Cancel(self.handle))
|
|
return False
|
|
|
|
def __get_query_state(self):
|
|
state, ok = self.__do_rpc(lambda : self.imp_service.get_state(self.handle))
|
|
if ok:
|
|
return state
|
|
else:
|
|
return self.query_state["EXCEPTION"]
|
|
|
|
def __do_rpc(self, rpc):
|
|
"""Executes the RPC lambda provided with some error checking. Returns
|
|
(rpc_result, True) if request was successful, (None, False) otherwise"""
|
|
if not self.connected:
|
|
print "Not connected (use CONNECT to establish a connection)"
|
|
return (None, False)
|
|
try:
|
|
return (rpc(), True)
|
|
except BeeswaxService.QueryNotFoundException, q:
|
|
print 'Error: Stale query handle'
|
|
# beeswaxException prints out the enture object, printing
|
|
# just the message a far more readable/helpful.
|
|
except BeeswaxService.BeeswaxException, b:
|
|
print "ERROR: %s" % (b.message,)
|
|
except TTransportException, e:
|
|
print "Error communicating with impalad: %s" % (e,)
|
|
self.connected = False
|
|
self.prompt = self.disconnected_prompt
|
|
except TApplicationException, t:
|
|
print "Application Exception : %s" % (t,)
|
|
except Exception, u:
|
|
print 'Unknown Exception : %s' % (u,)
|
|
return (None, False)
|
|
|
|
def do_explain(self, args):
|
|
query = BeeswaxService.Query()
|
|
# Args is all text except for 'explain', so no need to strip it out
|
|
query.query = args
|
|
query.configuration = self.__options_to_string_list()
|
|
print "Explain query: %s" % (query.query,)
|
|
(explanation, ok) = self.__do_rpc(lambda: self.imp_service.explain(query))
|
|
if ok:
|
|
print explanation.textual
|
|
return False
|
|
|
|
def do_refresh(self, args):
|
|
(_, ok) = self.__do_rpc(lambda: self.imp_service.ResetCatalog())
|
|
if ok:
|
|
print "Successfully refreshed catalog"
|
|
return False
|
|
|
|
def default(self, args):
|
|
print "Unrecognized command"
|
|
return False
|
|
|
|
def emptyline(self):
|
|
"""If an empty line is entered, do nothing"""
|
|
return False
|
|
|
|
def do_version(self, args):
|
|
"""Prints the Impala build version"""
|
|
print "Build version: %s" % VERSION_STRING
|
|
return False
|
|
|
|
def do_EOF(self, args):
|
|
"""Exit the shell on Ctrl^d"""
|
|
print 'Goodbye'
|
|
self.is_alive = False
|
|
return True
|
|
|
|
WELCOME_STRING = """Welcome to the Impala shell. Press TAB twice to see a list of \
|
|
available commands.
|
|
|
|
Copyright (c) 2012 Cloudera, Inc. All rights reserved.
|
|
(Build version: %s)""" % VERSION_STRING
|
|
|
|
if __name__ == "__main__":
|
|
parser = OptionParser()
|
|
parser.add_option("-i", "--impalad", dest="impalad", default=None,
|
|
help="<host:port> of impalad to connect to")
|
|
(options, args) = parser.parse_args()
|
|
|
|
intro = WELCOME_STRING
|
|
shell = ImpalaShell(options)
|
|
while shell.is_alive:
|
|
try:
|
|
shell.cmdloop(intro)
|
|
except KeyboardInterrupt:
|
|
intro = '\n'
|
|
pass
|