models.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524
  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 base64
  18. import datetime
  19. import logging
  20. import json
  21. from django.db import models
  22. from django.contrib.auth.models import User
  23. from django.contrib.contenttypes.fields import GenericRelation
  24. from django.core.urlresolvers import reverse
  25. from django.utils.translation import ugettext as _, ugettext_lazy as _t
  26. from enum import Enum
  27. from librdbms.server import dbms as librdbms_dbms
  28. from desktop.redaction import global_redaction_engine
  29. from desktop.lib.exceptions_renderable import PopupException
  30. from desktop.models import Document
  31. from TCLIService.ttypes import TSessionHandle, THandleIdentifier,\
  32. TOperationState, TOperationHandle, TOperationType
  33. from beeswax.design import HQLdesign
  34. LOG = logging.getLogger(__name__)
  35. QUERY_SUBMISSION_TIMEOUT = datetime.timedelta(0, 60 * 60) # 1 hour
  36. # Constants for DB fields, hue ini
  37. BEESWAX = 'beeswax'
  38. HIVE_SERVER2 = 'hiveserver2'
  39. QUERY_TYPES = (HQL, IMPALA, RDBMS, SPARK) = range(4)
  40. class QueryHistory(models.Model):
  41. """
  42. Holds metadata about all queries that have been executed.
  43. """
  44. STATE = Enum('submitted', 'running', 'available', 'failed', 'expired')
  45. SERVER_TYPE = ((BEESWAX, 'Beeswax'), (HIVE_SERVER2, 'Hive Server 2'),
  46. (librdbms_dbms.MYSQL, 'MySQL'), (librdbms_dbms.POSTGRESQL, 'PostgreSQL'),
  47. (librdbms_dbms.SQLITE, 'sqlite'), (librdbms_dbms.ORACLE, 'oracle'))
  48. owner = models.ForeignKey(User, db_index=True)
  49. query = models.TextField()
  50. last_state = models.IntegerField(db_index=True)
  51. has_results = models.BooleanField(default=False) # If true, this query will eventually return tabular results.
  52. submission_date = models.DateTimeField(auto_now_add=True)
  53. # In case of multi statements in a query, these are the id of the currently running statement
  54. server_id = models.CharField(max_length=1024, null=True) # Aka secret, only query in the "submitted" state is allowed to have no server_id
  55. server_guid = models.CharField(max_length=1024, null=True, default=None)
  56. statement_number = models.SmallIntegerField(default=0) # The index of the currently running statement
  57. operation_type = models.SmallIntegerField(null=True)
  58. modified_row_count = models.FloatField(null=True)
  59. log_context = models.CharField(max_length=1024, null=True)
  60. server_host = models.CharField(max_length=128, help_text=_('Host of the query server.'), default='')
  61. server_port = models.PositiveIntegerField(help_text=_('Port of the query server.'), default=10000)
  62. server_name = models.CharField(max_length=128, help_text=_('Name of the query server.'), default='')
  63. server_type = models.CharField(max_length=128, help_text=_('Type of the query server.'), default=BEESWAX, choices=SERVER_TYPE)
  64. query_type = models.SmallIntegerField(help_text=_('Type of the query.'), default=HQL, choices=((HQL, 'HQL'), (IMPALA, 'IMPALA')))
  65. design = models.ForeignKey('SavedQuery', to_field='id', null=True) # Some queries (like read/create table) don't have a design
  66. notify = models.BooleanField(default=False) # Notify on completion
  67. is_redacted = models.BooleanField(default=False)
  68. extra = models.TextField(default='{}') # Json fields for extra properties
  69. is_cleared = models.BooleanField(default=False)
  70. class Meta:
  71. ordering = ['-submission_date']
  72. @staticmethod
  73. def build(*args, **kwargs):
  74. return HiveServerQueryHistory(*args, **kwargs)
  75. def get_full_object(self):
  76. return HiveServerQueryHistory.objects.get(id=self.id)
  77. @staticmethod
  78. def get(id):
  79. return HiveServerQueryHistory.objects.get(id=id)
  80. @staticmethod
  81. def get_type_name(query_type):
  82. if query_type == IMPALA:
  83. return 'impala'
  84. elif query_type == RDBMS:
  85. return 'rdbms'
  86. elif query_type == SPARK:
  87. return 'spark'
  88. else:
  89. return 'beeswax'
  90. def get_query_server_config(self):
  91. from beeswax.server.dbms import get_query_server_config
  92. query_server = get_query_server_config(QueryHistory.get_type_name(self.query_type))
  93. query_server.update({
  94. 'server_name': self.server_name,
  95. # 'server_host': self.server_host, # Always use the live server configuration as the session is currently tied to the connection
  96. # 'server_port': int(self.server_port),
  97. 'server_type': self.server_type,
  98. })
  99. return query_server
  100. def get_current_statement(self):
  101. if self.design is not None:
  102. design = self.design.get_design()
  103. return design.get_query_statement(self.statement_number)
  104. else:
  105. return self.query
  106. def refresh_design(self, hql_query):
  107. # Refresh only HQL query part
  108. query = self.design.get_design()
  109. query.hql_query = hql_query
  110. self.design.data = query.dumps()
  111. self.query = hql_query
  112. def is_finished(self):
  113. is_statement_finished = not self.is_running()
  114. if self.design is not None:
  115. design = self.design.get_design()
  116. return is_statement_finished and self.statement_number + 1 == design.statement_count # Last statement
  117. else:
  118. return is_statement_finished
  119. def is_running(self):
  120. return self.last_state in (QueryHistory.STATE.running.index, QueryHistory.STATE.submitted.index)
  121. def is_success(self):
  122. return self.last_state in (QueryHistory.STATE.available.index,)
  123. def is_failure(self):
  124. return self.last_state in (QueryHistory.STATE.expired.index, QueryHistory.STATE.failed.index)
  125. def is_expired(self):
  126. return self.last_state in (QueryHistory.STATE.expired.index,)
  127. def set_to_running(self):
  128. self.last_state = QueryHistory.STATE.running.index
  129. def set_to_failed(self):
  130. self.last_state = QueryHistory.STATE.failed.index
  131. def set_to_available(self):
  132. self.last_state = QueryHistory.STATE.available.index
  133. def set_to_expired(self):
  134. self.last_state = QueryHistory.STATE.expired.index
  135. def save(self, *args, **kwargs):
  136. """
  137. Override `save` to optionally mask out the query from being saved to the
  138. database. This is because if the beeswax database contains sensitive
  139. information like personally identifiable information, that information
  140. could be leaked into the Hue database and logfiles.
  141. """
  142. if global_redaction_engine.is_enabled():
  143. redacted_query = global_redaction_engine.redact(self.query)
  144. if self.query != redacted_query:
  145. self.query = redacted_query
  146. self.is_redacted = True
  147. super(QueryHistory, self).save(*args, **kwargs)
  148. def update_extra(self, key, val):
  149. extra = json.loads(self.extra)
  150. extra[key] = val
  151. self.extra = json.dumps(extra)
  152. def get_extra(self, key):
  153. return json.loads(self.extra).get(key)
  154. def make_query_context(type, info):
  155. """
  156. ``type`` is one of "table" and "design", and ``info`` is the table name or design id.
  157. Returns a value suitable for GET param.
  158. """
  159. if type == 'table':
  160. return "%s:%s" % (type, info)
  161. elif type == 'design':
  162. # Use int() to validate that info is a number
  163. return "%s:%s" % (type, int(info))
  164. LOG.error("Invalid query context type: %s" % (type,))
  165. return '' # Empty string is safer than None
  166. class HiveServerQueryHistory(QueryHistory):
  167. # Map from (thrift) server state
  168. STATE_MAP = {
  169. TOperationState.INITIALIZED_STATE : QueryHistory.STATE.submitted,
  170. TOperationState.RUNNING_STATE : QueryHistory.STATE.running,
  171. TOperationState.FINISHED_STATE : QueryHistory.STATE.available,
  172. TOperationState.CANCELED_STATE : QueryHistory.STATE.failed,
  173. TOperationState.CLOSED_STATE : QueryHistory.STATE.expired,
  174. TOperationState.ERROR_STATE : QueryHistory.STATE.failed,
  175. TOperationState.UKNOWN_STATE : QueryHistory.STATE.failed,
  176. TOperationState.PENDING_STATE : QueryHistory.STATE.submitted,
  177. }
  178. node_type = HIVE_SERVER2
  179. class Meta:
  180. proxy = True
  181. def get_handle(self):
  182. secret, guid = HiveServerQueryHandle.get_decoded(self.server_id, self.server_guid)
  183. return HiveServerQueryHandle(secret=secret,
  184. guid=guid,
  185. has_result_set=self.has_results,
  186. operation_type=self.operation_type,
  187. modified_row_count=self.modified_row_count)
  188. def save_state(self, new_state):
  189. self.last_state = new_state.index
  190. self.save()
  191. @classmethod
  192. def is_canceled(self, res):
  193. return res.operationState in (TOperationState.CANCELED_STATE, TOperationState.CLOSED_STATE)
  194. class SavedQuery(models.Model):
  195. """
  196. Stores the query that people have save or submitted.
  197. Note that this used to be called QueryDesign. Any references to 'design'
  198. probably mean a SavedQuery.
  199. """
  200. DEFAULT_NEW_DESIGN_NAME = _('My saved query')
  201. AUTO_DESIGN_SUFFIX = _(' (new)')
  202. TYPES = QUERY_TYPES
  203. TYPES_MAPPING = {'beeswax': HQL, 'hql': HQL, 'impala': IMPALA, 'rdbms': RDBMS, 'spark': SPARK}
  204. type = models.IntegerField(null=False)
  205. owner = models.ForeignKey(User, db_index=True)
  206. # Data is a json of dictionary. See the beeswax.design module.
  207. data = models.TextField(max_length=65536)
  208. name = models.CharField(max_length=80)
  209. desc = models.TextField(max_length=1024)
  210. mtime = models.DateTimeField(auto_now=True)
  211. # An auto design is a place-holder for things users submit but not saved.
  212. # We still want to store it as a design to allow users to save them later.
  213. is_auto = models.BooleanField(default=False, db_index=True)
  214. is_trashed = models.BooleanField(default=False, db_index=True, verbose_name=_t('Is trashed'),
  215. help_text=_t('If this query is trashed.'))
  216. is_redacted = models.BooleanField(default=False)
  217. doc = GenericRelation(Document, related_query_name='hql_doc')
  218. class Meta:
  219. ordering = ['-mtime']
  220. def get_design(self):
  221. try:
  222. return HQLdesign.loads(self.data)
  223. except ValueError:
  224. # data is empty
  225. pass
  226. def clone(self, new_owner=None):
  227. if new_owner is None:
  228. new_owner = self.owner
  229. design = SavedQuery(type=self.type, owner=new_owner)
  230. design.data = self.data
  231. design.name = self.name
  232. design.desc = self.desc
  233. design.is_auto = self.is_auto
  234. return design
  235. @classmethod
  236. def create_empty(cls, app_name, owner, data):
  237. query_type = SavedQuery.TYPES_MAPPING[app_name]
  238. design = SavedQuery(owner=owner, type=query_type)
  239. design.name = SavedQuery.DEFAULT_NEW_DESIGN_NAME
  240. design.desc = ''
  241. if global_redaction_engine.is_enabled():
  242. design.data = global_redaction_engine.redact(data)
  243. else:
  244. design.data = data
  245. design.is_auto = True
  246. design.save()
  247. Document.objects.link(design, owner=design.owner, extra=design.type, name=design.name, description=design.desc)
  248. design.doc.get().add_to_history()
  249. return design
  250. @staticmethod
  251. def get(id, owner=None, type=None):
  252. """
  253. get(id, owner=None, type=None) -> SavedQuery object
  254. Checks that the owner and type match (when given).
  255. May raise PopupException (type/owner mismatch).
  256. May raise SavedQuery.DoesNotExist.
  257. """
  258. try:
  259. design = SavedQuery.objects.get(id=id)
  260. except SavedQuery.DoesNotExist, err:
  261. msg = _('Cannot retrieve query id %(id)s.') % {'id': id}
  262. raise err
  263. if owner is not None and design.owner != owner:
  264. msg = _('Query id %(id)s does not belong to user %(user)s.') % {'id': id, 'user': owner}
  265. LOG.error(msg)
  266. raise PopupException(msg)
  267. if type is not None and design.type != type:
  268. msg = _('Type mismatch for design id %(id)s (owner %(owner)s) - Expected %(expected_type)s, got %(real_type)s.') % \
  269. {'id': id, 'owner': owner, 'expected_type': design.type, 'real_type': type}
  270. LOG.error(msg)
  271. raise PopupException(msg)
  272. return design
  273. def __str__(self):
  274. return '%s %s' % (self.name, self.owner)
  275. def get_query_context(self):
  276. try:
  277. return make_query_context('design', self.id)
  278. except:
  279. LOG.exception('failed to make query context')
  280. return ""
  281. def get_absolute_url(self):
  282. return reverse(QueryHistory.get_type_name(self.type) + ':execute_design', kwargs={'design_id': self.id})
  283. def save(self, *args, **kwargs):
  284. """
  285. Override `save` to optionally mask out the query from being saved to the
  286. database. This is because if the beeswax database contains sensitive
  287. information like personally identifiable information, that information
  288. could be leaked into the Hue database and logfiles.
  289. """
  290. if global_redaction_engine.is_enabled():
  291. data = json.loads(self.data)
  292. try:
  293. query = data['query']['query']
  294. except KeyError:
  295. pass
  296. else:
  297. redacted_query = global_redaction_engine.redact(query)
  298. if query != redacted_query:
  299. data['query']['query'] = redacted_query
  300. self.is_redacted = True
  301. self.data = json.dumps(data)
  302. super(SavedQuery, self).save(*args, **kwargs)
  303. class SessionManager(models.Manager):
  304. def get_session(self, user, application='beeswax', filter_open=True):
  305. try:
  306. q = self.filter(owner=user, application=application).exclude(guid='').exclude(secret='')
  307. if filter_open:
  308. q = q.filter(status_code=0)
  309. return q.latest("last_used")
  310. except Session.DoesNotExist, e:
  311. return None
  312. def get_n_sessions(self, user, n, application='beeswax', filter_open=True):
  313. q = self.filter(owner=user, application=application).exclude(guid='').exclude(secret='')
  314. if filter_open:
  315. q = q.filter(status_code=0)
  316. return q.order_by("-last_used")[0:n]
  317. class Session(models.Model):
  318. """
  319. A sessions is bound to a user and an application (e.g. Bob with the Impala application).
  320. """
  321. owner = models.ForeignKey(User, db_index=True)
  322. status_code = models.PositiveSmallIntegerField() # ttypes.TStatusCode
  323. secret = models.TextField(max_length='100')
  324. guid = models.TextField(max_length='100')
  325. server_protocol_version = models.SmallIntegerField(default=0)
  326. last_used = models.DateTimeField(auto_now=True, db_index=True, verbose_name=_t('Last used'))
  327. application = models.CharField(max_length=128, help_text=_t('Application we communicate with.'), default='beeswax')
  328. properties = models.TextField(default='{}')
  329. objects = SessionManager()
  330. def get_handle(self):
  331. secret, guid = HiveServerQueryHandle.get_decoded(secret=self.secret, guid=self.guid)
  332. handle_id = THandleIdentifier(secret=secret, guid=guid)
  333. return TSessionHandle(sessionId=handle_id)
  334. def get_properties(self):
  335. return json.loads(self.properties) if self.properties else {}
  336. def get_formatted_properties(self):
  337. return [dict({'key': key, 'value': value}) for key, value in self.get_properties().items()]
  338. def __str__(self):
  339. return '%s %s' % (self.owner, self.last_used)
  340. class QueryHandle(object):
  341. def __init__(self, secret=None, guid=None, operation_type=None, has_result_set=None, modified_row_count=None, log_context=None, session_guid=None):
  342. self.secret = secret
  343. self.guid = guid
  344. self.operation_type = operation_type
  345. self.has_result_set = has_result_set
  346. self.modified_row_count = modified_row_count
  347. self.log_context = log_context
  348. def is_valid(self):
  349. return sum([bool(obj) for obj in [self.get()]]) > 0
  350. def __str__(self):
  351. return '%s %s' % (self.secret, self.guid)
  352. class HiveServerQueryHandle(QueryHandle):
  353. """
  354. QueryHandle for Hive Server 2.
  355. Store THandleIdentifier base64 encoded in order to be unicode compatible with Django.
  356. Also store session handle if provided.
  357. """
  358. def __init__(self, **kwargs):
  359. super(HiveServerQueryHandle, self).__init__(**kwargs)
  360. self.secret, self.guid = self.get_encoded()
  361. self.session_guid = kwargs.get('session_guid')
  362. def get(self):
  363. return self.secret, self.guid
  364. def get_rpc_handle(self):
  365. secret, guid = self.get_decoded(self.secret, self.guid)
  366. operation = getattr(TOperationType, TOperationType._NAMES_TO_VALUES.get(self.operation_type, 'EXECUTE_STATEMENT'))
  367. return TOperationHandle(operationId=THandleIdentifier(guid=guid, secret=secret),
  368. operationType=operation,
  369. hasResultSet=self.has_result_set,
  370. modifiedRowCount=self.modified_row_count)
  371. @classmethod
  372. def get_decoded(cls, secret, guid):
  373. return base64.decodestring(secret), base64.decodestring(guid)
  374. def get_encoded(self):
  375. return base64.encodestring(self.secret), base64.encodestring(self.guid)
  376. # Deprecated. Could be removed.
  377. class BeeswaxQueryHandle(QueryHandle):
  378. """
  379. QueryHandle for Beeswax.
  380. """
  381. def __init__(self, secret, has_result_set, log_context):
  382. super(BeeswaxQueryHandle, self).__init__(secret=secret,
  383. has_result_set=has_result_set,
  384. log_context=log_context)
  385. def get(self):
  386. return self.secret, None
  387. def get_rpc_handle(self):
  388. return BeeswaxdQueryHandle(id=self.secret, log_context=self.log_context)
  389. # TODO remove
  390. def get_encoded(self):
  391. return self.get(), None
  392. class MetaInstall(models.Model):
  393. """
  394. Metadata about the installation. Should have at most one row.
  395. """
  396. installed_example = models.BooleanField(default=False)
  397. @staticmethod
  398. def get():
  399. """
  400. MetaInstall.get() -> MetaInstall object
  401. It helps dealing with that this table has just one row.
  402. """
  403. try:
  404. return MetaInstall.objects.get(id=1)
  405. except MetaInstall.DoesNotExist:
  406. return MetaInstall(id=1)