hiveserver2.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567
  1. #!/usr/bin/env python
  2. # Licensed to Cloudera, Inc. under one
  3. # or more contributor license agreements. See the NOTICE file
  4. # distributed with this work for additional information
  5. # regarding copyright ownership. Cloudera, Inc. licenses this file
  6. # to you under the Apache License, Version 2.0 (the
  7. # "License"); you may not use this file except in compliance
  8. # with the License. You may obtain a copy of the License at
  9. #
  10. # http://www.apache.org/licenses/LICENSE-2.0
  11. #
  12. # Unless required by applicable law or agreed to in writing, software
  13. # distributed under the License is distributed on an "AS IS" BASIS,
  14. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  15. # See the License for the specific language governing permissions and
  16. # limitations under the License.
  17. import copy
  18. import logging
  19. import re
  20. import StringIO
  21. from django.core.urlresolvers import reverse
  22. from django.utils.translation import ugettext as _
  23. from desktop.lib.exceptions_renderable import PopupException
  24. from desktop.lib.i18n import force_unicode
  25. from desktop.models import DefaultConfiguration
  26. from notebook.connectors.base import Api, QueryError, QueryExpired
  27. LOG = logging.getLogger(__name__)
  28. try:
  29. from beeswax import data_export
  30. from beeswax.api import _autocomplete, _get_sample_data
  31. from beeswax.conf import CONFIG_WHITELIST as hive_settings
  32. from beeswax.data_export import upload
  33. from beeswax.design import hql_query, strip_trailing_semicolon, split_statements
  34. from beeswax import conf as beeswax_conf
  35. from beeswax.models import QUERY_TYPES, HiveServerQueryHandle, HiveServerQueryHistory, QueryHistory, Session
  36. from beeswax.server import dbms
  37. from beeswax.server.dbms import get_query_server_config, QueryServerException
  38. from beeswax.views import _parse_out_hadoop_jobs
  39. except ImportError, e:
  40. LOG.exception('Hive and HiveServer2 interfaces are not enabled')
  41. try:
  42. from impala.conf import CONFIG_WHITELIST as impala_settings
  43. from impala import views # Force checking if Impala is enabled
  44. except ImportError, e:
  45. LOG.warn("Impala app is not enabled")
  46. impala_settings = None
  47. DEFAULT_HIVE_ENGINE = 'mr'
  48. def query_error_handler(func):
  49. def decorator(*args, **kwargs):
  50. try:
  51. return func(*args, **kwargs)
  52. except QueryServerException, e:
  53. message = force_unicode(str(e))
  54. if 'Invalid query handle' in message or 'Invalid OperationHandle' in message:
  55. raise QueryExpired(e)
  56. else:
  57. raise QueryError(message)
  58. return decorator
  59. class HiveConfiguration(object):
  60. APP_NAME = 'hive'
  61. PROPERTIES = [
  62. {
  63. "multiple": True,
  64. "defaultValue": [],
  65. "value": [],
  66. "nice_name": _("Files"),
  67. "key": "files",
  68. "help_text": _("Add one or more files, jars, or archives to the list of resources."),
  69. "type": "hdfs-files"
  70. }, {
  71. "multiple": True,
  72. "defaultValue": [],
  73. "value": [],
  74. "nice_name": _("Functions"),
  75. "key": "functions",
  76. "help_text": _("Add one or more registered UDFs (requires function name and fully-qualified class name)."),
  77. "type": "functions"
  78. }, {
  79. "multiple": True,
  80. "defaultValue": [],
  81. "value": [],
  82. "nice_name": _("Settings"),
  83. "key": "settings",
  84. "help_text": _("Hive and Hadoop configuration properties."),
  85. "type": "settings",
  86. "options": [config.lower() for config in hive_settings.get()]
  87. }
  88. ]
  89. class ImpalaConfiguration(object):
  90. APP_NAME = 'impala'
  91. PROPERTIES = [
  92. {
  93. "multiple": True,
  94. "defaultValue": [],
  95. "value": [],
  96. "nice_name": _("Settings"),
  97. "key": "settings",
  98. "help_text": _("Impala configuration properties."),
  99. "type": "settings",
  100. "options": [config.lower() for config in impala_settings.get()] if impala_settings is not None else []
  101. }
  102. ]
  103. class HS2Api(Api):
  104. @staticmethod
  105. def get_properties(lang='hive'):
  106. return ImpalaConfiguration.PROPERTIES if lang == 'impala' else HiveConfiguration.PROPERTIES
  107. @query_error_handler
  108. def create_session(self, lang='hive', properties=None):
  109. application = 'beeswax' if lang == 'hive' else lang
  110. session = Session.objects.get_session(self.user, application=application)
  111. if session is None:
  112. session = dbms.get(self.user, query_server=get_query_server_config(name=lang)).open_session(self.user)
  113. if not properties:
  114. config = DefaultConfiguration.objects.get_configuration_for_user(app=lang, user=self.user)
  115. if config is not None:
  116. properties = config.properties_list
  117. else:
  118. properties = self.get_properties(lang)
  119. return {
  120. 'type': lang,
  121. 'id': session.id,
  122. 'properties': properties
  123. }
  124. @query_error_handler
  125. def close_session(self, session):
  126. app_name = session.get('type')
  127. session_id = session.get('id')
  128. query_server = get_query_server_config(name=app_name)
  129. response = {'status': -1, 'message': ''}
  130. try:
  131. filters = {'id': session_id, 'application': query_server['server_name']}
  132. if not self.user.is_superuser:
  133. filters['owner'] = self.user
  134. session = Session.objects.get(**filters)
  135. except Session.DoesNotExist:
  136. response['message'] = _('Session does not exist or you do not have permissions to close the session.')
  137. if session:
  138. session = dbms.get(self.user, query_server).close_session(session)
  139. response['status'] = 0
  140. response['message'] = _('Session successfully closed.')
  141. response['session'] = {'id': session_id, 'application': session.application, 'status': session.status_code}
  142. return response
  143. @query_error_handler
  144. def execute(self, notebook, snippet):
  145. db = self._get_db(snippet)
  146. statement = self._get_current_statement(db, snippet)
  147. session = self._get_session(notebook, snippet['type'])
  148. query = self._prepare_hql_query(snippet, statement['statement'], session)
  149. try:
  150. db.use(query.database)
  151. handle = db.client.query(query)
  152. except QueryServerException, ex:
  153. raise QueryError(ex.message, handle=statement)
  154. # All good
  155. server_id, server_guid = handle.get()
  156. response = {
  157. 'secret': server_id,
  158. 'guid': server_guid,
  159. 'operation_type': handle.operation_type,
  160. 'has_result_set': handle.has_result_set,
  161. 'modified_row_count': handle.modified_row_count,
  162. 'log_context': handle.log_context,
  163. }
  164. response.update(statement)
  165. return response
  166. @query_error_handler
  167. def check_status(self, notebook, snippet):
  168. response = {}
  169. db = self._get_db(snippet)
  170. handle = self._get_handle(snippet)
  171. operation = db.get_operation_status(handle)
  172. status = HiveServerQueryHistory.STATE_MAP[operation.operationState]
  173. if status.index in (QueryHistory.STATE.failed.index, QueryHistory.STATE.expired.index):
  174. raise QueryError(operation.errorMessage)
  175. response['status'] = 'running' if status.index in (QueryHistory.STATE.running.index, QueryHistory.STATE.submitted.index) else 'available'
  176. return response
  177. @query_error_handler
  178. def fetch_result(self, notebook, snippet, rows, start_over):
  179. db = self._get_db(snippet)
  180. handle = self._get_handle(snippet)
  181. results = db.fetch(handle, start_over=start_over, rows=rows)
  182. # No escaping...
  183. return {
  184. 'has_more': results.has_more,
  185. 'data': results.rows(),
  186. 'meta': [{
  187. 'name': column.name,
  188. 'type': column.type,
  189. 'comment': column.comment
  190. } for column in results.data_table.cols()],
  191. 'type': 'table'
  192. }
  193. @query_error_handler
  194. def fetch_result_metadata(self):
  195. pass
  196. @query_error_handler
  197. def cancel(self, notebook, snippet):
  198. db = self._get_db(snippet)
  199. handle = self._get_handle(snippet)
  200. db.cancel_operation(handle)
  201. return {'status': 0}
  202. @query_error_handler
  203. def get_log(self, notebook, snippet, startFrom=None, size=None):
  204. db = self._get_db(snippet)
  205. handle = self._get_handle(snippet)
  206. return db.get_log(handle, start_over=startFrom == 0)
  207. @query_error_handler
  208. def close_statement(self, snippet):
  209. if snippet['type'] == 'impala':
  210. from impala import conf as impala_conf
  211. if (snippet['type'] == 'hive' and beeswax_conf.CLOSE_QUERIES.get()) or (snippet['type'] == 'impala' and impala_conf.CLOSE_QUERIES.get()):
  212. db = self._get_db(snippet)
  213. handle = self._get_handle(snippet)
  214. db.close_operation(handle)
  215. return {'status': 0}
  216. else:
  217. return {'status': -1} # skipped
  218. @query_error_handler
  219. def download(self, notebook, snippet, format):
  220. try:
  221. db = self._get_db(snippet)
  222. handle = self._get_handle(snippet)
  223. return data_export.download(handle, format, db)
  224. except Exception, e:
  225. LOG.exception('error downloading notebook')
  226. if not hasattr(e, 'message') or not e.message:
  227. message = e
  228. else:
  229. message = e.message
  230. raise PopupException(message, detail='')
  231. @query_error_handler
  232. def progress(self, snippet, logs):
  233. if snippet['type'] == 'hive':
  234. match = re.search('Total jobs = (\d+)', logs, re.MULTILINE)
  235. total = int(match.group(1)) if match else 1
  236. started = logs.count('Starting Job')
  237. ended = logs.count('Ended Job')
  238. progress = int((started + ended) * 100 / (total * 2))
  239. return max(progress, 5) # Return 5% progress as a minimum
  240. elif snippet['type'] == 'impala':
  241. match = re.search('(\d+)% Complete', logs, re.MULTILINE)
  242. return int(match.group(1)) if match else 0
  243. else:
  244. return 50
  245. @query_error_handler
  246. def get_jobs(self, notebook, snippet, logs):
  247. jobs = []
  248. if snippet['type'] == 'hive':
  249. engine = self._get_hive_execution_engine(notebook, snippet)
  250. job_ids = _parse_out_hadoop_jobs(logs, engine=engine)
  251. jobs = [{
  252. 'name': job_id,
  253. 'url': reverse('jobbrowser.views.single_job', kwargs={'job': job_id})
  254. } for job_id in job_ids]
  255. return jobs
  256. @query_error_handler
  257. def autocomplete(self, snippet, database=None, table=None, column=None, nested=None):
  258. db = self._get_db(snippet)
  259. return _autocomplete(db, database, table, column, nested)
  260. @query_error_handler
  261. def get_sample_data(self, snippet, database=None, table=None, column=None):
  262. db = self._get_db(snippet)
  263. return _get_sample_data(db, database, table, column)
  264. @query_error_handler
  265. def explain(self, notebook, snippet):
  266. db = self._get_db(snippet)
  267. response = self._get_current_statement(db, snippet)
  268. session = self._get_session(notebook, snippet['type'])
  269. query = self._prepare_hql_query(snippet, response.pop('statement'), session)
  270. explanation = db.explain(query)
  271. return {
  272. 'status': 0,
  273. 'explanation': explanation.textual,
  274. 'statement': query.get_query_statement(0),
  275. }
  276. @query_error_handler
  277. def export_data_as_hdfs_file(self, snippet, target_file, overwrite):
  278. db = self._get_db(snippet)
  279. handle = self._get_handle(snippet)
  280. upload(target_file, handle, self.request.user, db, self.request.fs)
  281. return '/filebrowser/view=%s' % target_file
  282. def export_data_as_table(self, notebook, snippet, destination):
  283. db = self._get_db(snippet)
  284. response = self._get_current_statement(db, snippet)
  285. session = self._get_session(notebook, snippet['type'])
  286. query = self._prepare_hql_query(snippet, response.pop('statement'), session)
  287. if not query.hql_query.strip().lower().startswith('select'):
  288. raise Exception(_('Only SELECT statements can be saved. Provided statement: %(query)s') % {'query': query.hql_query})
  289. database = snippet.get('database') or 'default'
  290. table = destination
  291. if '.' in table:
  292. database, table = table.split('.', 1)
  293. db.use(query.database)
  294. hql = 'CREATE TABLE `%s`.`%s` AS %s' % (database, table, query.hql_query)
  295. success_url = reverse('metastore:describe_table', kwargs={'database': database, 'table': table})
  296. return hql, success_url
  297. def export_large_data_to_hdfs(self, notebook, snippet, destination):
  298. db = self._get_db(snippet)
  299. response = self._get_current_statement(db, snippet)
  300. session = self._get_session(notebook, snippet['type'])
  301. query = self._prepare_hql_query(snippet, response.pop('statement'), session)
  302. if not query.hql_query.strip().lower().startswith('select'):
  303. raise Exception(_('Only SELECT statements can be saved. Provided statement: %(query)s') % {'query': query.hql_query})
  304. db.use(query.database)
  305. hql = "INSERT OVERWRITE DIRECTORY '%s' %s" % (destination, query.hql_query)
  306. success_url = '/filebrowser/view=%s' % destination
  307. return hql, success_url
  308. def upgrade_properties(self, lang='hive', properties=None):
  309. upgraded_properties = copy.deepcopy(self.get_properties(lang))
  310. # Check that current properties is a list of dictionary objects with 'key' and 'value' keys
  311. if not isinstance(properties, list) or \
  312. not all(isinstance(prop, dict) for prop in properties) or \
  313. not all('key' in prop for prop in properties) or not all('value' in prop for prop in properties):
  314. LOG.warn('Current properties are not formatted correctly, will replace with defaults.')
  315. return upgraded_properties
  316. valid_props_dict = dict((prop["key"], prop) for prop in upgraded_properties)
  317. curr_props_dict = dict((prop['key'], prop) for prop in properties)
  318. # Upgrade based on valid properties as needed
  319. if set(valid_props_dict.keys()) != set(curr_props_dict.keys()):
  320. settings = next((prop for prop in upgraded_properties if prop['key'] == 'settings'), None)
  321. if settings is not None and isinstance(properties, list):
  322. settings['value'] = properties
  323. else: # No upgrade needed so return existing properties
  324. upgraded_properties = properties
  325. return upgraded_properties
  326. def _get_session(self, notebook, type='hive'):
  327. session = next((session for session in notebook['sessions'] if session['type'] == type), None)
  328. return session
  329. def _get_hive_execution_engine(self, notebook, snippet):
  330. # Get hive.execution.engine from snippet properties, if none, then get from session
  331. properties = snippet['properties']
  332. settings = properties.get('settings', [])
  333. if not settings:
  334. session = self._get_session(notebook, 'hive')
  335. if not session:
  336. raise Exception(_('Cannot get jobs, failed to find active HS2 session for user: %s') % self.user.username)
  337. properties = session['properties']
  338. settings = next((prop['value'] for prop in properties if prop['key'] == 'settings'), None)
  339. if settings:
  340. engine = next((setting['value'] for setting in settings if setting['key'] == 'hive.execution.engine'), DEFAULT_HIVE_ENGINE)
  341. else:
  342. engine = DEFAULT_HIVE_ENGINE
  343. return engine
  344. def _get_statements(self, hql_query):
  345. hql_query = strip_trailing_semicolon(hql_query)
  346. hql_query_sio = StringIO.StringIO(hql_query)
  347. statements = []
  348. for (start_row, start_col), (end_row, end_col), statement in split_statements(hql_query_sio.read()):
  349. statements.append({
  350. 'start': {
  351. 'row': start_row,
  352. 'column': start_col
  353. },
  354. 'end': {
  355. 'row': end_row,
  356. 'column': end_col
  357. },
  358. 'statement': strip_trailing_semicolon(statement.strip())
  359. })
  360. return statements
  361. def _get_current_statement(self, db, snippet):
  362. # Multiquery, if not first statement or arrived to the last query
  363. statement_id = snippet['result']['handle'].get('statement_id', 0)
  364. statements_count = snippet['result']['handle'].get('statements_count', 1)
  365. if snippet['result']['handle'].get('has_more_statements'):
  366. try:
  367. handle = self._get_handle(snippet)
  368. db.close_operation(handle) # Close all the time past multi queries
  369. except:
  370. LOG.warn('Could not close previous multiquery query')
  371. statement_id += 1
  372. else:
  373. statement_id = 0
  374. statements = self._get_statements(snippet['statement'])
  375. resp = {
  376. 'statement_id': statement_id,
  377. 'has_more_statements': statement_id < len(statements) - 1,
  378. 'statements_count': len(statements)
  379. }
  380. if statements_count != len(statements):
  381. statement_id = 0
  382. resp.update(statements[statement_id])
  383. return resp
  384. def _prepare_hql_query(self, snippet, statement, session):
  385. settings = snippet['properties'].get('settings', None)
  386. file_resources = snippet['properties'].get('files', None)
  387. functions = snippet['properties'].get('functions', None)
  388. properties = session['properties']
  389. # Get properties from session if not defined in snippet
  390. if not settings:
  391. settings = next((prop['value'] for prop in properties if prop['key'] == 'settings'), None)
  392. if not file_resources:
  393. file_resources = next((prop['value'] for prop in properties if prop['key'] == 'files'), None)
  394. if not functions:
  395. functions = next((prop['value'] for prop in properties if prop['key'] == 'functions'), None)
  396. database = snippet.get('database') or 'default'
  397. return hql_query(
  398. statement,
  399. query_type=QUERY_TYPES[0],
  400. settings=settings,
  401. file_resources=file_resources,
  402. functions=functions,
  403. database=database
  404. )
  405. def get_select_star_query(self, snippet, database, table):
  406. db = self._get_db(snippet)
  407. table = db.get_table(database, table)
  408. return db.get_select_star_query(database, table)
  409. def _get_handle(self, snippet):
  410. snippet['result']['handle']['secret'], snippet['result']['handle']['guid'] = HiveServerQueryHandle.get_decoded(snippet['result']['handle']['secret'], snippet['result']['handle']['guid'])
  411. for key in snippet['result']['handle'].keys():
  412. if key not in ('log_context', 'secret', 'has_result_set', 'operation_type', 'modified_row_count', 'guid'):
  413. snippet['result']['handle'].pop(key)
  414. return HiveServerQueryHandle(**snippet['result']['handle'])
  415. def _get_db(self, snippet):
  416. if snippet['type'] == 'hive':
  417. name = 'beeswax'
  418. elif snippet['type'] == 'impala':
  419. name = 'impala'
  420. else:
  421. name = 'spark-sql'
  422. return dbms.get(self.user, query_server=get_query_server_config(name=name))