optimizer_client.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414
  1. #!/usr/bin/env python
  2. # -- coding: utf-8 --
  3. # Licensed to Cloudera, Inc. under one
  4. # or more contributor license agreements. See the NOTICE file
  5. # distributed with this work for additional information
  6. # regarding copyright ownership. Cloudera, Inc. licenses this file
  7. # to you under the Apache License, Version 2.0 (the
  8. # "License"); you may not use this file except in compliance
  9. # with the License. You may obtain a copy of the License at
  10. #
  11. # http://www.apache.org/licenses/LICENSE-2.0
  12. #
  13. # Unless required by applicable law or agreed to in writing, software
  14. # distributed under the License is distributed on an "AS IS" BASIS,
  15. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  16. # See the License for the specific language governing permissions and
  17. # limitations under the License.
  18. from builtins import object
  19. import json
  20. import logging
  21. import os
  22. import time
  23. import uuid
  24. from tempfile import NamedTemporaryFile
  25. from django.core.cache import cache
  26. from django.utils.functional import wraps
  27. from django.utils.translation import ugettext as _
  28. from desktop.auth.backend import is_admin
  29. from desktop.lib.exceptions_renderable import PopupException
  30. from desktop.lib import export_csvxls
  31. from desktop.lib.i18n import smart_unicode
  32. from desktop.lib.rest.http_client import RestException
  33. from libsentry.sentry_site import get_hive_sentry_provider
  34. from libsentry.privilege_checker import get_checker, MissingSentryPrivilegeException
  35. from metadata.conf import OPTIMIZER, get_optimizer_url
  36. LOG = logging.getLogger(__name__)
  37. _JSON_CONTENT_TYPE = 'application/json'
  38. OPTIMIZER_TENANT_ID_CACHE_KEY = 'navopt-tenant-id'
  39. try:
  40. from navoptapi.api_lib import ApiLib
  41. except Exception as e:
  42. LOG.warn('NavOpt module is not installed: %s' % e)
  43. class NavOptException(Exception):
  44. def __init__(self, message=None):
  45. self.message = message or _('No error message, please check the logs.')
  46. def __str__(self):
  47. return str(self.message)
  48. def __unicode__(self):
  49. return smart_unicode(self.message)
  50. def check_privileges(view_func):
  51. def decorate(*args, **kwargs):
  52. if OPTIMIZER.APPLY_SENTRY_PERMISSIONS.get():
  53. checker = get_checker(user=args[0].user)
  54. action = 'SELECT'
  55. objects = []
  56. if kwargs.get('db_tables'):
  57. for db_table in kwargs['db_tables']:
  58. objects.append({'server': get_hive_sentry_provider(), 'db': _get_table_name(db_table)['database'], 'table': _get_table_name(db_table)['table']})
  59. else:
  60. objects = [{'server': get_hive_sentry_provider()}]
  61. if kwargs.get('database_name'):
  62. objects[0]['db'] = kwargs['database_name']
  63. if kwargs.get('table_name'):
  64. objects[0]['table'] = kwargs['table_name']
  65. filtered = list(checker.filter_objects(objects, action))
  66. if len(filtered) != len(objects):
  67. raise MissingSentryPrivilegeException({'pre_filtering': objects, 'post_filtering': filtered, 'diff': len(objects) - len(filtered)})
  68. return view_func(*args, **kwargs)
  69. return wraps(view_func)(decorate)
  70. class OptimizerApi(object):
  71. def __init__(self, user, api_url=None, auth_key=None, auth_key_secret=None, tenant_id=None):
  72. self.user = user
  73. self._api_url = (api_url or get_optimizer_url()).strip('/')
  74. self._auth_key = auth_key if auth_key else OPTIMIZER.AUTH_KEY_ID.get()
  75. self._auth_key_secret = auth_key_secret if auth_key_secret else (OPTIMIZER.AUTH_KEY_SECRET.get() and OPTIMIZER.AUTH_KEY_SECRET.get().replace('\\n', '\n'))
  76. self._api = ApiLib("navopt", self._api_url, self._auth_key, self._auth_key_secret)
  77. self._tenant_id = tenant_id if tenant_id else _get_tenant_id(self) # Aka "workload"
  78. def _call(self, *kwargs):
  79. start_time = time.time()
  80. resp = self._api.call_api(*kwargs)
  81. data = resp.json()
  82. if resp.headers.get('x-altus-request-id'):
  83. LOG.info('%s %s in %dms: %s' % (self.user, resp.headers['x-altus-request-id'], (time.time() - start_time) * 1000, kwargs))
  84. if data.get('code') == 'UNKNOWN':
  85. raise NavOptException(data.get('message'))
  86. elif data.get('errorMsg'):
  87. raise NavOptException(data.get('errorMsg'))
  88. else:
  89. return data
  90. def get_tenant(self, cluster_id='default'):
  91. return self._call('getTenant', {'clusterId' : cluster_id})
  92. def upload(self, data, data_type='queries', source_platform='generic', workload_id=None):
  93. if data_type in ('table_stats', 'cols_stats'):
  94. data_suffix = '.json'
  95. if data_type == 'table_stats':
  96. extra_parameters = {'fileType': 'TABLE_STATS'}
  97. else:
  98. extra_parameters = {'fileType': 'COLUMN_STATS'}
  99. else:
  100. data_suffix = '.csv'
  101. extra_parameters = {
  102. 'fileType': 'QUERY',
  103. 'colDelim': ',',
  104. 'rowDelim': '\n',
  105. "headerFields": [
  106. {"count": 0, "name": "SQL_ID", "coltype": "SQL_ID", "use": True, "tag": ""},
  107. {"count": 0, "name": "ELAPSED_TIME", "coltype": "NONE", "use": True, "tag": ""},
  108. {"count": 0, "name": "SQL_FULLTEXT", "coltype": "SQL_QUERY", "use": True, "tag": ""},
  109. {"count": 0, "name": "DATABASE", "coltype": "NONE", "use": True, "tag": "DATABASE"}
  110. ],
  111. }
  112. f_queries_path = NamedTemporaryFile(suffix=data_suffix)
  113. f_queries_path.close() # Reopened as real file below to work well with the command
  114. try:
  115. f_queries = open(f_queries_path.name, 'w+')
  116. try:
  117. # Queries
  118. if data_suffix == '.csv':
  119. content_generator = OptimizerQueryDataAdapter(data)
  120. queries_csv = export_csvxls.create_generator(content_generator, 'csv')
  121. for row in queries_csv:
  122. f_queries.write(row)
  123. LOG.debug(row[:1000])
  124. else:
  125. # Table, column stats
  126. f_queries.write(json.dumps(data))
  127. LOG.debug(json.dumps(data[:10]))
  128. finally:
  129. f_queries.close()
  130. parameters = {
  131. 'tenant' : self._tenant_id,
  132. 'fileLocation': f_queries.name,
  133. 'sourcePlatform': source_platform,
  134. }
  135. parameters.update(extra_parameters)
  136. response = self._api.call_api('upload', parameters)
  137. status = json.loads(response)
  138. status['count'] = len(data)
  139. return status
  140. except RestException as e:
  141. raise PopupException(e, title=_('Error while accessing Optimizer'))
  142. finally:
  143. os.remove(f_queries_path.name)
  144. def upload_status(self, workload_id):
  145. return self._call('uploadStatus', {'tenant' : self._tenant_id, 'workloadId': workload_id})
  146. # Sentry permissions work bottom to top.
  147. # @check_privileges
  148. def top_tables(self, workfloadId=None, database_name='default', page_size=1000, startingToken=None):
  149. data = self._call('getTopTables', {'tenant' : self._tenant_id, 'dbName': database_name.lower(), 'pageSize': page_size, 'startingToken': startingToken})
  150. if OPTIMIZER.APPLY_SENTRY_PERMISSIONS.get():
  151. checker = get_checker(user=self.user)
  152. action = 'SELECT'
  153. def getkey(table):
  154. names = _get_table_name(table['name'])
  155. return {'server': get_hive_sentry_provider(), 'db': names['database'], 'table': names['table']}
  156. data['results'] = list(checker.filter_objects(data['results'], action, key=getkey))
  157. return data
  158. @check_privileges
  159. def table_details(self, database_name, table_name, page_size=100, startingToken=None):
  160. return self._call('getTablesDetail', {'tenant' : self._tenant_id, 'dbName': database_name.lower(), 'tableName': table_name.lower(), 'pageSize': page_size, 'startingToken': startingToken})
  161. def query_compatibility(self, source_platform, target_platform, query, page_size=100, startingToken=None):
  162. return self._call('getQueryCompatible', {'tenant' : self._tenant_id, 'query': query, 'sourcePlatform': source_platform, 'targetPlatform': target_platform, 'startingToken': startingToken})
  163. def query_risk(self, query, source_platform, db_name, page_size=100, startingToken=None):
  164. response = self._call('getQueryRisk', {
  165. 'tenant' : self._tenant_id,
  166. 'query': _clean_query(query),
  167. 'dbName': db_name,
  168. 'sourcePlatform': source_platform,
  169. 'pageSize': page_size,
  170. 'startingToken': startingToken
  171. })
  172. hints = response.get(source_platform + 'Risk', {})
  173. if hints and hints == [{"riskTables": [], "riskAnalysis": "", "riskId": 0, "risk": "low", "riskRecommendation": ""}]:
  174. hints = []
  175. return {
  176. 'hints': hints,
  177. 'noStats': response.get('noStats', []),
  178. 'noDDL': response.get('noDDL', []),
  179. }
  180. def similar_queries(self, source_platform, query, page_size=100, startingToken=None):
  181. if is_admin(self.user):
  182. return self._call('getSimilarQueries', {'tenant' : self._tenant_id, 'sourcePlatform': source_platform, 'query': query, 'pageSize': page_size, 'startingToken': startingToken})
  183. else:
  184. raise PopupException(_('Call not supported'))
  185. @check_privileges
  186. def top_filters(self, db_tables=None, page_size=100, startingToken=None):
  187. args = {
  188. 'tenant' : self._tenant_id,
  189. 'pageSize': page_size,
  190. 'startingToken': startingToken
  191. }
  192. if db_tables:
  193. args['dbTableList'] = [db_table.lower() for db_table in db_tables]
  194. results = self._call('getTopFilters', args)
  195. if OPTIMIZER.APPLY_SENTRY_PERMISSIONS.get():
  196. filtered_filters = []
  197. for result in results['results']:
  198. cols = [_get_table_name(col['columnName']) for col in result["popularValues"][0]["group"]]
  199. if len(cols) == len(list(_secure_results(cols, self.user))):
  200. filtered_filters.append(result)
  201. results['results'] = filtered_filters
  202. return results
  203. @check_privileges
  204. def top_aggs(self, db_tables=None, page_size=100, startingToken=None):
  205. args = {
  206. 'tenant' : self._tenant_id,
  207. 'pageSize': page_size,
  208. 'startingToken': startingToken
  209. }
  210. if db_tables:
  211. args['dbTableList'] = [db_table.lower() for db_table in db_tables]
  212. results = self._call('getTopAggs', args)
  213. if OPTIMIZER.APPLY_SENTRY_PERMISSIONS.get():
  214. checker = get_checker(user=self.user)
  215. action = 'SELECT'
  216. def getkey(table):
  217. names = table['aggregateInfo'][0]
  218. names['server'] = get_hive_sentry_provider()
  219. return names
  220. results['results'] = list(checker.filter_objects(results['results'], action, key=getkey))
  221. return results
  222. @check_privileges
  223. def top_columns(self, db_tables=None, page_size=100, startingToken=None):
  224. args = {
  225. 'tenant' : self._tenant_id,
  226. 'pageSize': page_size,
  227. 'startingToken': startingToken
  228. }
  229. if db_tables:
  230. args['dbTableList'] = [db_table.lower() for db_table in db_tables]
  231. results = self._call('getTopColumns', args)
  232. if OPTIMIZER.APPLY_SENTRY_PERMISSIONS.get():
  233. for section in ['orderbyColumns', 'selectColumns', 'filterColumns', 'joinColumns', 'groupbyColumns']:
  234. results[section] = list(_secure_results(results[section], self.user))
  235. return results
  236. @check_privileges
  237. def top_joins(self, db_tables=None, page_size=100, startingToken=None):
  238. args = {
  239. 'tenant' : self._tenant_id,
  240. 'pageSize': page_size,
  241. 'startingToken': startingToken
  242. }
  243. if db_tables:
  244. args['dbTableList'] = [db_table.lower() for db_table in db_tables]
  245. results = self._call('getTopJoins', args)
  246. if OPTIMIZER.APPLY_SENTRY_PERMISSIONS.get():
  247. filtered_joins = []
  248. for result in results['results']:
  249. cols = [_get_table_name(col) for col in result["joinCols"][0]["columns"]]
  250. if len(cols) == len(list(_secure_results(cols, self.user))):
  251. filtered_joins.append(result)
  252. results['results'] = filtered_joins
  253. return results
  254. def top_databases(self, page_size=100, startingToken=None):
  255. args = {
  256. 'tenant' : self._tenant_id,
  257. 'pageSize': page_size,
  258. 'startingToken': startingToken
  259. }
  260. data = self._call('getTopDatabases', args)
  261. if OPTIMIZER.APPLY_SENTRY_PERMISSIONS.get():
  262. data['results'] = list(_secure_results(data['results'], self.user))
  263. return data
  264. def OptimizerQueryDataAdapter(data):
  265. headers = ['SQL_ID', 'ELAPSED_TIME', 'SQL_FULLTEXT', 'DATABASE']
  266. if data and len(data[0]) == 4:
  267. rows = data
  268. else:
  269. rows = ([str(uuid.uuid4()), 0.0, q, 'default'] for q in data)
  270. yield headers, rows
  271. def _get_table_name(path):
  272. column = None
  273. if path.count('.') == 1:
  274. database, table = path.split('.', 1)
  275. elif path.count('.') == 2:
  276. database, table, column = path.split('.', 2)
  277. else:
  278. database, table = 'default', path
  279. name = {'database': database, 'table': table}
  280. if column:
  281. name['column'] = column
  282. return name
  283. def _secure_results(results, user, action='SELECT'):
  284. if OPTIMIZER.APPLY_SENTRY_PERMISSIONS.get():
  285. checker = get_checker(user=user)
  286. def getkey(result):
  287. key = {'server': get_hive_sentry_provider()}
  288. if 'dbName' in result:
  289. key['db'] = result['dbName']
  290. elif 'database' in result:
  291. key['db'] = result['database']
  292. if 'tableName' in result:
  293. key['table'] = result['tableName']
  294. elif 'table' in result:
  295. key['table'] = result['table']
  296. if 'columnName' in result:
  297. key['column'] = result['columnName']
  298. elif 'column' in result:
  299. key['column'] = result['column']
  300. return key
  301. return checker.filter_objects(results, action, key=getkey)
  302. else:
  303. return results
  304. def _clean_query(query):
  305. return ' '.join([line for line in query.strip().splitlines() if not line.strip().startswith('--')])
  306. def _get_tenant_id(api):
  307. tenant_id = OPTIMIZER.TENANT_ID.get() or cache.get(OPTIMIZER_TENANT_ID_CACHE_KEY)
  308. if not tenant_id:
  309. tenant = api.get_tenant(cluster_id=OPTIMIZER.CLUSTER_ID.get())
  310. if tenant.get('tenant'):
  311. tenant_id = tenant['tenant']
  312. else:
  313. raise PopupException(_('Could not get tenant id from cluster id %s: %s') % (OPTIMIZER.CLUSTER_ID.get(), tenant))
  314. cache.set(OPTIMIZER_TENANT_ID_CACHE_KEY, tenant_id, 60 * 60 * 24 * 30)
  315. return tenant_id