hiveserver2.py 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305
  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 logging
  18. import re
  19. from django.core.urlresolvers import reverse
  20. from django.utils.translation import ugettext as _
  21. from desktop.lib.exceptions_renderable import PopupException
  22. from desktop.lib.i18n import force_unicode
  23. from notebook.connectors.base import Api, QueryError, QueryExpired
  24. LOG = logging.getLogger(__name__)
  25. try:
  26. from beeswax import data_export
  27. from beeswax.api import _autocomplete
  28. from beeswax.design import hql_query, strip_trailing_semicolon, split_statements
  29. from beeswax import conf as beeswax_conf
  30. from beeswax.models import QUERY_TYPES, HiveServerQueryHandle, HiveServerQueryHistory, QueryHistory, Session
  31. from beeswax.server import dbms
  32. from beeswax.server.dbms import get_query_server_config, QueryServerException
  33. from beeswax.views import _parse_out_hadoop_jobs
  34. except ImportError, e:
  35. LOG.exception('Hive and HiveServer2 interfaces are not enabled')
  36. def query_error_handler(func):
  37. def decorator(*args, **kwargs):
  38. try:
  39. return func(*args, **kwargs)
  40. except QueryServerException, e:
  41. message = force_unicode(str(e))
  42. if 'Invalid query handle' in message or 'Invalid OperationHandle' in message:
  43. raise QueryExpired(e)
  44. else:
  45. raise QueryError(message)
  46. return decorator
  47. class HS2Api(Api):
  48. @query_error_handler
  49. def create_session(self, lang='hive', properties=None):
  50. application = 'beeswax' if lang == 'hive' else lang
  51. session = Session.objects.get_session(self.user, application=application)
  52. if session is None:
  53. session = dbms.get(self.user, query_server=get_query_server_config(name=lang)).open_session(self.user)
  54. return {
  55. 'type': lang,
  56. 'id': session.id,
  57. 'properties': session.get_formatted_properties()
  58. }
  59. @query_error_handler
  60. def close_session(self, session):
  61. app_name = session.get('type')
  62. session_id = session.get('id')
  63. query_server = get_query_server_config(name=app_name)
  64. response = {'status': -1, 'message': ''}
  65. try:
  66. filters = {'id': session_id, 'application': query_server['server_name']}
  67. if not self.user.is_superuser:
  68. filters['owner'] = self.user
  69. session = Session.objects.get(**filters)
  70. except Session.DoesNotExist:
  71. response['message'] = _('Session does not exist or you do not have permissions to close the session.')
  72. if session:
  73. session = dbms.get(self.user, query_server).close_session(session)
  74. response['status'] = 0
  75. response['message'] = _('Session successfully closed.')
  76. response['session'] = {'id': session_id, 'application': session.application, 'status': session.status_code}
  77. return response
  78. @query_error_handler
  79. def execute(self, notebook, snippet):
  80. db = self._get_db(snippet)
  81. # Multiquery, if not first statement or arrived to the last query
  82. statement_id = snippet['result']['handle'].get('statement_id', 0)
  83. if snippet['result']['handle'].get('has_more_statements'):
  84. try:
  85. handle = self._get_handle(snippet)
  86. db.close_operation(handle) # Close all the time past multi queries
  87. except:
  88. LOG.warn('Could not close previous multiquery query')
  89. statement_id += 1
  90. else:
  91. statement_id = 0
  92. statements = self._get_statements(snippet['statement'])
  93. statement = statements[statement_id]
  94. settings = snippet['properties'].get('settings', None)
  95. file_resources = snippet['properties'].get('files', None)
  96. functions = snippet['properties'].get('functions', None)
  97. if settings:
  98. settings = [{'key': s.rsplit('=', 1)[0], 'value': s.rsplit('=', 1)[1]} for s in settings] # TODO integrate with new binding
  99. if file_resources:
  100. file_resources = [{'type': 'JAR', 'path': f} for f in file_resources] # TODO support FILE ARCHIVE
  101. if functions:
  102. functions = [{'name': f.rsplit(' ', 1)[0], 'class_name': f.rsplit(' ', 1)[1]} for f in functions] # TODO protect for index out of bounds
  103. query = hql_query(statement, query_type=QUERY_TYPES[0], settings=settings, file_resources=file_resources, functions=functions)
  104. try:
  105. handle = db.client.query(query)
  106. except QueryServerException, ex:
  107. raise QueryError(ex.message)
  108. # All good
  109. server_id, server_guid = handle.get()
  110. return {
  111. 'secret': server_id,
  112. 'guid': server_guid,
  113. 'operation_type': handle.operation_type,
  114. 'has_result_set': handle.has_result_set,
  115. 'modified_row_count': handle.modified_row_count,
  116. 'log_context': handle.log_context,
  117. 'statement_id': statement_id,
  118. 'has_more_statements': statement_id < len(statements) - 1
  119. }
  120. def _get_statements(self, hql_query):
  121. hql_query = strip_trailing_semicolon(hql_query)
  122. return [strip_trailing_semicolon(statement.strip()) for statement in split_statements(hql_query)]
  123. @query_error_handler
  124. def check_status(self, notebook, snippet):
  125. response = {}
  126. db = self._get_db(snippet)
  127. handle = self._get_handle(snippet)
  128. operation = db.get_operation_status(handle)
  129. status = HiveServerQueryHistory.STATE_MAP[operation.operationState]
  130. if status.index in (QueryHistory.STATE.failed.index, QueryHistory.STATE.expired.index):
  131. raise QueryError(operation.errorMessage)
  132. response['status'] = 'running' if status.index in (QueryHistory.STATE.running.index, QueryHistory.STATE.submitted.index) else 'available'
  133. return response
  134. @query_error_handler
  135. def fetch_result(self, notebook, snippet, rows, start_over):
  136. db = self._get_db(snippet)
  137. handle = self._get_handle(snippet)
  138. results = db.fetch(handle, start_over=start_over, rows=rows)
  139. # No escaping...
  140. return {
  141. 'has_more': results.has_more,
  142. 'data': results.rows(),
  143. 'meta': [{
  144. 'name': column.name,
  145. 'type': column.type,
  146. 'comment': column.comment
  147. } for column in results.data_table.cols()],
  148. 'type': 'table'
  149. }
  150. @query_error_handler
  151. def fetch_result_metadata(self):
  152. pass
  153. @query_error_handler
  154. def cancel(self, notebook, snippet):
  155. db = self._get_db(snippet)
  156. handle = self._get_handle(snippet)
  157. db.cancel_operation(handle)
  158. return {'status': 0}
  159. @query_error_handler
  160. def get_log(self, notebook, snippet, startFrom=None, size=None):
  161. db = self._get_db(snippet)
  162. handle = self._get_handle(snippet)
  163. return db.get_log(handle, start_over=startFrom == 0)
  164. @query_error_handler
  165. def close_statement(self, snippet):
  166. if snippet['type'] == 'impala':
  167. from impala import conf as impala_conf
  168. if (snippet['type'] == 'hive' and beeswax_conf.CLOSE_QUERIES.get()) or (snippet['type'] == 'impala' and impala_conf.CLOSE_QUERIES.get()):
  169. db = self._get_db(snippet)
  170. handle = self._get_handle(snippet)
  171. db.close_operation(handle)
  172. return {'status': 0}
  173. else:
  174. return {'status': -1} # skipped
  175. @query_error_handler
  176. def download(self, notebook, snippet, format):
  177. try:
  178. db = self._get_db(snippet)
  179. handle = self._get_handle(snippet)
  180. return data_export.download(handle, format, db)
  181. except Exception, e:
  182. LOG.exception('error downloading notebook')
  183. if not hasattr(e, 'message') or not e.message:
  184. message = e
  185. else:
  186. message = e.message
  187. raise PopupException(message, detail='')
  188. @query_error_handler
  189. def progress(self, snippet, logs):
  190. if snippet['type'] == 'hive':
  191. match = re.search('Total jobs = (\d+)', logs, re.MULTILINE)
  192. total = int(match.group(1)) if match else 1
  193. started = logs.count('Starting Job')
  194. ended = logs.count('Ended Job')
  195. progress = int((started + ended) * 100 / (total * 2))
  196. return max(progress, 5) # Return 5% progress as a minimum
  197. elif snippet['type'] == 'impala':
  198. match = re.search('(\d+)% Complete', logs, re.MULTILINE)
  199. return int(match.group(1)) if match else 0
  200. else:
  201. return 50
  202. @query_error_handler
  203. def get_jobs(self, notebook, snippet, logs):
  204. job_ids = _parse_out_hadoop_jobs(logs)
  205. jobs = [{
  206. 'name': job_id,
  207. 'url': reverse('jobbrowser.views.single_job', kwargs={'job': job_id})
  208. } for job_id in job_ids]
  209. return jobs
  210. @query_error_handler
  211. def autocomplete(self, snippet, database=None, table=None, column=None, nested=None):
  212. db = self._get_db(snippet)
  213. return _autocomplete(db, database, table, column, nested)
  214. def get_select_star_query(self, snippet, database, table):
  215. db = self._get_db(snippet)
  216. table = db.get_table(database, table)
  217. return db.get_select_star_query(database, table)
  218. def _get_handle(self, snippet):
  219. snippet['result']['handle']['secret'], snippet['result']['handle']['guid'] = HiveServerQueryHandle.get_decoded(snippet['result']['handle']['secret'], snippet['result']['handle']['guid'])
  220. snippet['result']['handle'].pop('statement_id')
  221. snippet['result']['handle'].pop('has_more_statements')
  222. return HiveServerQueryHandle(**snippet['result']['handle'])
  223. def _get_db(self, snippet):
  224. if snippet['type'] == 'hive':
  225. name = 'beeswax'
  226. elif snippet['type'] == 'impala':
  227. name = 'impala'
  228. else:
  229. name = 'spark-sql'
  230. return dbms.get(self.user, query_server=get_query_server_config(name=name))