api.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946
  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. from future import standard_library
  18. standard_library.install_aliases()
  19. import json
  20. import logging
  21. import sqlparse
  22. import sys
  23. from django.urls import reverse
  24. from django.db.models import Q
  25. from django.utils.translation import ugettext as _
  26. from django.views.decorators.http import require_GET, require_POST
  27. import opentracing.tracer
  28. from azure.abfs.__init__ import abfspath
  29. from desktop.conf import TASK_SERVER
  30. from desktop.lib.i18n import smart_str
  31. from desktop.lib.django_util import JsonResponse
  32. from desktop.models import Document2, Document, __paginate, _get_gist_document
  33. from indexer.file_format import HiveFormat
  34. from indexer.fields import Field
  35. from notebook.connectors.base import Notebook, QueryExpired, SessionExpired, QueryError, _get_snippet_name
  36. from notebook.connectors.hiveserver2 import HS2Api
  37. from notebook.decorators import api_error_handler, check_document_access_permission, check_document_modify_permission
  38. from notebook.models import escape_rows, make_notebook, upgrade_session_properties, get_api
  39. if sys.version_info[0] > 2:
  40. from urllib.parse import unquote as urllib_unquote
  41. else:
  42. from urllib import unquote as urllib_unquote
  43. LOG = logging.getLogger(__name__)
  44. DEFAULT_HISTORY_NAME = ''
  45. @require_POST
  46. @api_error_handler
  47. def create_notebook(request):
  48. response = {'status': -1}
  49. editor_type = request.POST.get('type', 'notebook')
  50. gist_id = request.POST.get('gist')
  51. directory_uuid = request.POST.get('directory_uuid')
  52. if gist_id:
  53. gist_doc = _get_gist_document(uuid=gist_id)
  54. statement = json.loads(gist_doc.data)['statement']
  55. editor = make_notebook(
  56. name='',
  57. description='',
  58. editor_type=editor_type,
  59. statement=statement,
  60. is_presentation_mode=True
  61. )
  62. else:
  63. editor = Notebook()
  64. data = editor.get_data()
  65. if editor_type != 'notebook':
  66. data['name'] = ''
  67. data['type'] = 'query-%s' % editor_type # TODO: Add handling for non-SQL types
  68. data['directoryUuid'] = directory_uuid
  69. editor.data = json.dumps(data)
  70. response['notebook'] = editor.get_data()
  71. response['status'] = 0
  72. return JsonResponse(response)
  73. @require_POST
  74. @check_document_access_permission
  75. @api_error_handler
  76. def create_session(request):
  77. response = {'status': -1}
  78. session = json.loads(request.POST.get('session', '{}'))
  79. properties = session.get('properties', [])
  80. response['session'] = get_api(request, session).create_session(lang=session['type'], properties=properties)
  81. response['status'] = 0
  82. return JsonResponse(response)
  83. @require_POST
  84. @check_document_access_permission
  85. @api_error_handler
  86. def close_session(request):
  87. response = {'status': -1}
  88. session = json.loads(request.POST.get('session', '{}'))
  89. response['session'] = get_api(request, {'type': session['type']}).close_session(session=session)
  90. response['status'] = 0
  91. return JsonResponse(response)
  92. def _execute_notebook(request, notebook, snippet):
  93. response = {'status': -1}
  94. result = None
  95. history = None
  96. historify = (notebook['type'] != 'notebook' or snippet.get('wasBatchExecuted')) and not notebook.get('skipHistorify')
  97. try:
  98. try:
  99. session = notebook.get('sessions') and notebook['sessions'][0] # Session reference for snippet execution without persisting it
  100. if historify:
  101. history = _historify(notebook, request.user)
  102. notebook = Notebook(document=history).get_data()
  103. interpreter = get_api(request, snippet)
  104. if snippet.get('interface') == 'sqlalchemy':
  105. interpreter.options['session'] = session
  106. with opentracing.tracer.start_span('interpreter') as span:
  107. response['handle'] = interpreter.execute(notebook, snippet)
  108. # Retrieve and remove the result from the handle
  109. if response['handle'].get('sync'):
  110. result = response['handle'].pop('result')
  111. finally:
  112. if historify:
  113. _snippet = [s for s in notebook['snippets'] if s['id'] == snippet['id']][0]
  114. if 'handle' in response: # No failure
  115. if 'result' not in _snippet: # Editor v2
  116. _snippet['result'] = {}
  117. _snippet['result']['handle'] = response['handle']
  118. _snippet['result']['statements_count'] = response['handle'].get('statements_count', 1)
  119. _snippet['result']['statement_id'] = response['handle'].get('statement_id', 0)
  120. _snippet['result']['handle']['statement'] = response['handle'].get('statement', snippet['statement']).strip() # For non HS2, as non multi query yet
  121. else:
  122. _snippet['status'] = 'failed'
  123. if history: # If _historify failed, history will be None. If we get Atomic block exception, something underneath interpreter.execute() crashed and is not handled.
  124. history.update_data(notebook)
  125. history.save()
  126. response['history_id'] = history.id
  127. response['history_uuid'] = history.uuid
  128. if notebook['isSaved']: # Keep track of history of saved queries
  129. response['history_parent_uuid'] = history.dependencies.filter(type__startswith='query-').latest('last_modified').uuid
  130. except QueryError as ex: # We inject the history information from _historify() to the failed queries
  131. if response.get('history_id'):
  132. ex.extra['history_id'] = response['history_id']
  133. if response.get('history_uuid'):
  134. ex.extra['history_uuid'] = response['history_uuid']
  135. if response.get('history_parent_uuid'):
  136. ex.extra['history_parent_uuid'] = response['history_parent_uuid']
  137. raise ex
  138. # Inject and HTML escape results
  139. if result is not None:
  140. response['result'] = result
  141. response['result']['data'] = escape_rows(result['data'])
  142. response['status'] = 0
  143. return response
  144. @require_POST
  145. @check_document_access_permission
  146. @api_error_handler
  147. def execute(request, engine=None):
  148. notebook = json.loads(request.POST.get('notebook', '{}'))
  149. snippet = json.loads(request.POST.get('snippet', '{}'))
  150. with opentracing.tracer.start_span('notebook-execute') as span:
  151. span.set_tag('user-id', request.user.username)
  152. response = _execute_notebook(request, notebook, snippet)
  153. span.set_tag(
  154. 'query-id',
  155. response['handle']['guid'] if response.get('handle') and response['handle'].get('guid') else None
  156. )
  157. return JsonResponse(response)
  158. @require_POST
  159. @check_document_access_permission
  160. @api_error_handler
  161. def check_status(request):
  162. response = {'status': -1}
  163. operation_id = request.POST.get('operationId')
  164. notebook = json.loads(request.POST.get('notebook', '{}'))
  165. snippet = json.loads(request.POST.get('snippet', '{}'))
  166. if operation_id or not snippet: # To unify with _get_snippet
  167. nb_doc = Document2.objects.get_by_uuid(user=request.user, uuid=operation_id or notebook['id'])
  168. notebook = Notebook(document=nb_doc).get_data() # Used below
  169. snippet = notebook['snippets'][0]
  170. try:
  171. with opentracing.tracer.start_span('notebook-check_status') as span:
  172. span.set_tag('user-id', request.user.username)
  173. span.set_tag(
  174. 'query-id',
  175. snippet['result']['handle']['guid'] if snippet['result'].get('handle') and snippet['result']['handle'].get('guid') else None
  176. )
  177. response['query_status'] = get_api(request, snippet).check_status(notebook, snippet)
  178. response['status'] = 0
  179. except SessionExpired:
  180. response['status'] = 'expired'
  181. raise
  182. except QueryExpired:
  183. response['status'] = 'expired'
  184. raise
  185. finally:
  186. if response['status'] == 0 and snippet['status'] != response['query_status']:
  187. status = response['query_status']['status']
  188. elif response['status'] == 'expired':
  189. status = 'expired'
  190. else:
  191. status = 'failed'
  192. if notebook['type'].startswith('query') or notebook.get('isManaged'):
  193. nb_doc = Document2.objects.get(id=notebook['id'])
  194. if nb_doc.can_write(request.user):
  195. nb = Notebook(document=nb_doc).get_data()
  196. if status != nb['snippets'][0]['status']:
  197. nb['snippets'][0]['status'] = status
  198. nb_doc.update_data(nb)
  199. nb_doc.save()
  200. return JsonResponse(response)
  201. @require_POST
  202. @check_document_access_permission
  203. @api_error_handler
  204. def fetch_result_data(request):
  205. response = {'status': -1}
  206. operation_id = request.POST.get('operationId')
  207. notebook = json.loads(request.POST.get('notebook', '{}'))
  208. snippet = json.loads(request.POST.get('snippet', '{}'))
  209. rows = json.loads(request.POST.get('rows', '100'))
  210. start_over = json.loads(request.POST.get('startOver', 'false'))
  211. snippet = _get_snippet(request.user, notebook, snippet, operation_id)
  212. with opentracing.tracer.start_span('notebook-fetch_result_data') as span:
  213. response['result'] = get_api(request, snippet).fetch_result(notebook, snippet, rows, start_over)
  214. span.set_tag('user-id', request.user.username)
  215. span.set_tag(
  216. 'query-id',
  217. snippet['result']['handle']['guid'] if snippet['result'].get('handle') and snippet['result']['handle'].get('guid') else None
  218. )
  219. # Materialize and HTML escape results
  220. if response['result'].get('data') and response['result'].get('type') == 'table' and not response['result'].get('isEscaped'):
  221. response['result']['data'] = escape_rows(response['result']['data'])
  222. response['result']['isEscaped'] = True
  223. response['status'] = 0
  224. return JsonResponse(response)
  225. @require_POST
  226. @check_document_access_permission
  227. @api_error_handler
  228. def fetch_result_metadata(request):
  229. response = {'status': -1}
  230. operation_id = request.POST.get('operationId')
  231. notebook = json.loads(request.POST.get('notebook', '{}'))
  232. snippet = json.loads(request.POST.get('snippet', '{}'))
  233. snippet = _get_snippet(request.user, notebook, snippet, operation_id)
  234. with opentracing.tracer.start_span('notebook-fetch_result_metadata') as span:
  235. response['result'] = get_api(request, snippet).fetch_result_metadata(notebook, snippet)
  236. span.set_tag('user-id', request.user.username)
  237. span.set_tag(
  238. 'query-id',
  239. snippet['result']['handle']['guid'] if snippet['result'].get('handle') and snippet['result']['handle'].get('guid') else None
  240. )
  241. response['status'] = 0
  242. return JsonResponse(response)
  243. @require_POST
  244. @check_document_access_permission
  245. @api_error_handler
  246. def fetch_result_size(request):
  247. response = {'status': -1}
  248. operation_id = request.POST.get('operationId')
  249. notebook = json.loads(request.POST.get('notebook', '{}'))
  250. snippet = json.loads(request.POST.get('snippet', '{}'))
  251. snippet = _get_snippet(request.user, notebook, snippet, operation_id)
  252. with opentracing.tracer.start_span('notebook-fetch_result_size') as span:
  253. response['result'] = get_api(request, snippet).fetch_result_size(notebook, snippet)
  254. span.set_tag('user-id', request.user.username)
  255. span.set_tag(
  256. 'query-id',
  257. snippet['result']['handle']['guid'] if snippet['result'].get('handle') and snippet['result']['handle'].get('guid') else None
  258. )
  259. response['status'] = 0
  260. return JsonResponse(response)
  261. @require_POST
  262. @check_document_access_permission
  263. @api_error_handler
  264. def cancel_statement(request):
  265. response = {'status': -1}
  266. notebook = json.loads(request.POST.get('notebook', '{}'))
  267. snippet = None
  268. operation_id = request.POST.get('operationId') or notebook['uuid']
  269. snippet = _get_snippet(request.user, notebook, snippet, operation_id)
  270. with opentracing.tracer.start_span('notebook-cancel_statement') as span:
  271. response['result'] = get_api(request, snippet).cancel(notebook, snippet)
  272. span.set_tag('user-id', request.user.username)
  273. span.set_tag(
  274. 'query-id',
  275. snippet['result']['handle']['guid'] if snippet['result'].get('handle') and snippet['result']['handle'].get('guid') else None
  276. )
  277. response['status'] = 0
  278. return JsonResponse(response)
  279. @require_POST
  280. @check_document_access_permission
  281. @api_error_handler
  282. def get_logs(request):
  283. response = {'status': -1}
  284. operation_id = request.POST.get('operationId')
  285. notebook = json.loads(request.POST.get('notebook', '{}'))
  286. snippet = json.loads(request.POST.get('snippet', '{}'))
  287. startFrom = request.POST.get('from')
  288. startFrom = int(startFrom) if startFrom else None
  289. size = request.POST.get('size')
  290. size = int(size) if size else None
  291. full_log = smart_str(request.POST.get('full_log', ''))
  292. snippet = _get_snippet(request.user, notebook, snippet, operation_id)
  293. db = get_api(request, snippet)
  294. with opentracing.tracer.start_span('notebook-get_logs') as span:
  295. logs = smart_str(db.get_log(notebook, snippet, startFrom=startFrom, size=size))
  296. span.set_tag('user-id', request.user.username)
  297. span.set_tag(
  298. 'query-id',
  299. snippet['result']['handle']['guid'] if snippet['result'].get('handle') and snippet['result']['handle'].get('guid') else None
  300. )
  301. full_log += logs
  302. jobs = db.get_jobs(notebook, snippet, full_log)
  303. response['logs'] = logs.strip()
  304. response['progress'] = min(db.progress(notebook, snippet, logs=full_log), 99) if snippet['status'] != 'available' and snippet['status'] != 'success' else 100
  305. response['jobs'] = jobs
  306. response['isFullLogs'] = db.get_log_is_full_log(notebook, snippet)
  307. response['status'] = 0
  308. return JsonResponse(response)
  309. def _save_notebook(notebook, user):
  310. notebook_type = notebook.get('type', 'notebook')
  311. save_as = False
  312. if notebook.get('parentSavedQueryUuid'): # We save into the original saved query, not into the query history
  313. notebook_doc = Document2.objects.get_by_uuid(user=user, uuid=notebook['parentSavedQueryUuid'])
  314. elif notebook.get('id'):
  315. notebook_doc = Document2.objects.get(id=notebook['id'])
  316. else:
  317. notebook_doc = Document2.objects.create(name=notebook['name'], uuid=notebook['uuid'], type=notebook_type, owner=user)
  318. Document.objects.link(notebook_doc, owner=notebook_doc.owner, name=notebook_doc.name, description=notebook_doc.description, extra=notebook_type)
  319. save_as = True
  320. if notebook.get('directoryUuid'):
  321. notebook_doc.parent_directory = Document2.objects.get_by_uuid(user=user, uuid=notebook.get('directoryUuid'), perm_type='write')
  322. else:
  323. notebook_doc.parent_directory = Document2.objects.get_home_directory(user)
  324. notebook['isSaved'] = True
  325. notebook['isHistory'] = False
  326. notebook['id'] = notebook_doc.id
  327. _clear_sessions(notebook)
  328. notebook_doc1 = notebook_doc._get_doc1(doc2_type=notebook_type)
  329. notebook_doc.update_data(notebook)
  330. notebook_doc.search = _get_statement(notebook)
  331. notebook_doc.name = notebook_doc1.name = notebook['name']
  332. notebook_doc.description = notebook_doc1.description = notebook['description']
  333. notebook_doc.save()
  334. notebook_doc1.save()
  335. return notebook_doc, save_as
  336. @api_error_handler
  337. @require_POST
  338. @check_document_modify_permission()
  339. def save_notebook(request):
  340. response = {'status': -1}
  341. notebook = json.loads(request.POST.get('notebook', '{}'))
  342. notebook_doc, save_as = _save_notebook(notebook, request.user)
  343. response['status'] = 0
  344. response['save_as'] = save_as
  345. response.update(notebook_doc.to_dict())
  346. response['message'] = request.POST.get('editorMode') == 'true' and _('Query saved successfully') or _('Notebook saved successfully')
  347. return JsonResponse(response)
  348. def _clear_sessions(notebook):
  349. notebook['sessions'] = [_s for _s in notebook['sessions'] if _s['type'] in ('scala', 'spark', 'pyspark', 'sparkr', 'r')]
  350. def _historify(notebook, user):
  351. query_type = notebook['type']
  352. name = notebook['name'] if (notebook['name'] and notebook['name'].strip() != '') else DEFAULT_HISTORY_NAME
  353. is_managed = notebook.get('isManaged') == True # Prevents None
  354. if is_managed and Document2.objects.filter(uuid=notebook['uuid']).exists():
  355. history_doc = Document2.objects.get(uuid=notebook['uuid'])
  356. else:
  357. history_doc = Document2.objects.create(
  358. name=name,
  359. type=query_type,
  360. owner=user,
  361. is_history=True,
  362. is_managed=is_managed
  363. )
  364. # Link history of saved query
  365. if notebook['isSaved']:
  366. parent_doc = Document2.objects.get(uuid=notebook.get('parentSavedQueryUuid') or notebook['uuid']) # From previous history query or initial saved query
  367. notebook['parentSavedQueryUuid'] = parent_doc.uuid
  368. history_doc.dependencies.add(parent_doc)
  369. if not is_managed:
  370. Document.objects.link(
  371. history_doc,
  372. name=history_doc.name,
  373. owner=history_doc.owner,
  374. description=history_doc.description,
  375. extra=query_type
  376. )
  377. notebook['uuid'] = history_doc.uuid
  378. _clear_sessions(notebook)
  379. history_doc.update_data(notebook)
  380. history_doc.search = _get_statement(notebook)
  381. history_doc.save()
  382. return history_doc
  383. def _get_statement(notebook):
  384. if notebook['snippets'] and len(notebook['snippets']) > 0:
  385. return Notebook.statement_with_variables(notebook['snippets'][0])
  386. return ''
  387. @require_GET
  388. @api_error_handler
  389. @check_document_access_permission
  390. def get_history(request):
  391. response = {'status': -1}
  392. doc_type = request.GET.get('doc_type')
  393. doc_text = request.GET.get('doc_text')
  394. page = min(int(request.GET.get('page', 1)), 100)
  395. limit = min(int(request.GET.get('limit', 50)), 100)
  396. is_notification_manager = request.GET.get('is_notification_manager', 'false') == 'true'
  397. if is_notification_manager:
  398. docs = Document2.objects.get_tasks_history(user=request.user)
  399. else:
  400. docs = Document2.objects.get_history(doc_type='query-%s' % doc_type, user=request.user)
  401. if doc_text:
  402. docs = docs.filter(Q(name__icontains=doc_text) | Q(description__icontains=doc_text) | Q(search__icontains=doc_text))
  403. # Paginate
  404. docs = docs.order_by('-last_modified')
  405. response['count'] = docs.count()
  406. docs = __paginate(page, limit, queryset=docs)['documents']
  407. history = []
  408. for doc in docs:
  409. notebook = Notebook(document=doc).get_data()
  410. if 'snippets' in notebook:
  411. statement = notebook['description'] if is_notification_manager else _get_statement(notebook)
  412. history.append({
  413. 'name': doc.name,
  414. 'id': doc.id,
  415. 'uuid': doc.uuid,
  416. 'type': doc.type,
  417. 'data': {
  418. 'statement': statement[:1001] if statement else '',
  419. 'lastExecuted': notebook['snippets'][0].get('lastExecuted', -1),
  420. 'status': notebook['snippets'][0]['status'],
  421. 'parentSavedQueryUuid': notebook.get('parentSavedQueryUuid', '')
  422. } if notebook['snippets'] else {},
  423. 'absoluteUrl': doc.get_absolute_url(),
  424. })
  425. else:
  426. LOG.error('Incomplete History Notebook: %s' % notebook)
  427. response['history'] = sorted(history, key=lambda row: row['data']['lastExecuted'], reverse=True)
  428. response['message'] = _('History fetched')
  429. response['status'] = 0
  430. return JsonResponse(response)
  431. @require_POST
  432. @api_error_handler
  433. @check_document_modify_permission()
  434. def clear_history(request):
  435. response = {'status': -1}
  436. notebook = json.loads(request.POST.get('notebook', '{}'))
  437. doc_type = request.POST.get('doc_type')
  438. is_notification_manager = request.POST.get('is_notification_manager', 'false') == 'true'
  439. if is_notification_manager:
  440. history = Document2.objects.get_tasks_history(user=request.user)
  441. else:
  442. history = Document2.objects.get_history(doc_type='query-%s' % doc_type, user=request.user)
  443. response['updated'] = history.delete()
  444. response['message'] = _('History cleared !')
  445. response['status'] = 0
  446. return JsonResponse(response)
  447. @require_GET
  448. @check_document_access_permission
  449. def open_notebook(request):
  450. response = {'status': -1}
  451. notebook_id = request.GET.get('notebook')
  452. notebook = Notebook(document=Document2.objects.get(id=notebook_id))
  453. notebook = upgrade_session_properties(request, notebook)
  454. response['status'] = 0
  455. response['notebook'] = notebook.get_json()
  456. response['message'] = _('Notebook loaded successfully')
  457. @require_POST
  458. @check_document_access_permission
  459. def close_notebook(request):
  460. response = {'status': -1, 'result': []}
  461. notebook = json.loads(request.POST.get('notebook', '{}'))
  462. for session in [_s for _s in notebook['sessions'] if _s['type'] in ('scala', 'spark', 'pyspark', 'sparkr', 'r')]:
  463. try:
  464. response['result'].append(get_api(request, session).close_session(session))
  465. except QueryExpired:
  466. pass
  467. except Exception as e:
  468. LOG.exception('Error closing session %s' % str(e))
  469. for snippet in [_s for _s in notebook['snippets'] if _s['type'] in ('hive', 'impala')]:
  470. try:
  471. if snippet['status'] != 'running':
  472. response['result'].append(get_api(request, snippet).close_statement(notebook, snippet))
  473. else:
  474. LOG.info('Not closing SQL snippet as still running.')
  475. except QueryExpired:
  476. pass
  477. except Exception as e:
  478. LOG.exception('Error closing statement %s' % str(e))
  479. response['status'] = 0
  480. response['message'] = _('Notebook closed successfully')
  481. return JsonResponse(response)
  482. @require_POST
  483. @check_document_access_permission
  484. def close_statement(request):
  485. response = {'status': -1}
  486. notebook = json.loads(request.POST.get('notebook', '{}'))
  487. snippet = None
  488. operation_id = request.POST.get('operationId') or notebook['uuid']
  489. snippet = _get_snippet(request.user, notebook, snippet, operation_id)
  490. try:
  491. with opentracing.tracer.start_span('notebook-close_statement') as span:
  492. response['result'] = get_api(request, snippet).close_statement(notebook, snippet)
  493. span.set_tag('user-id', request.user.username)
  494. span.set_tag(
  495. 'query-id',
  496. snippet['result']['handle']['guid'] if snippet['result'].get('handle') and snippet['result']['handle'].get('guid') else None
  497. )
  498. except QueryExpired:
  499. pass
  500. response['status'] = 0
  501. response['message'] = _('Statement closed !')
  502. return JsonResponse(response)
  503. @require_POST
  504. @check_document_access_permission
  505. @api_error_handler
  506. def autocomplete(request, server=None, database=None, table=None, column=None, nested=None):
  507. response = {'status': -1}
  508. # Passed by check_document_access_permission but unused by APIs
  509. notebook = json.loads(request.POST.get('notebook', '{}'))
  510. snippet = json.loads(request.POST.get('snippet', '{}'))
  511. try:
  512. autocomplete_data = get_api(request, snippet).autocomplete(snippet, database, table, column, nested)
  513. response.update(autocomplete_data)
  514. except QueryExpired:
  515. pass
  516. response['status'] = 0
  517. return JsonResponse(response)
  518. @require_POST
  519. @check_document_access_permission
  520. @api_error_handler
  521. def get_sample_data(request, server=None, database=None, table=None, column=None):
  522. response = {'status': -1}
  523. # Passed by check_document_access_permission but unused by APIs
  524. notebook = json.loads(request.POST.get('notebook', '{}'))
  525. snippet = json.loads(request.POST.get('snippet', '{}'))
  526. async = json.loads(request.POST.get('async', 'false'))
  527. operation = json.loads(request.POST.get('operation', '"default"'))
  528. sample_data = get_api(request, snippet).get_sample_data(snippet, database, table, column, async=async, operation=operation)
  529. response.update(sample_data)
  530. response['status'] = 0
  531. return JsonResponse(response)
  532. @require_POST
  533. @check_document_access_permission
  534. @api_error_handler
  535. def explain(request):
  536. response = {'status': -1}
  537. notebook = json.loads(request.POST.get('notebook', '{}'))
  538. snippet = json.loads(request.POST.get('snippet', '{}'))
  539. response = get_api(request, snippet).explain(notebook, snippet)
  540. return JsonResponse(response)
  541. @require_POST
  542. @api_error_handler
  543. def format(request):
  544. response = {'status': 0}
  545. statements = request.POST.get('statements', '')
  546. response['formatted_statements'] = sqlparse.format(statements, reindent=True, keyword_case='upper') # SQL only currently
  547. return JsonResponse(response)
  548. @require_POST
  549. @check_document_access_permission
  550. @api_error_handler
  551. def export_result(request):
  552. response = {'status': -1, 'message': _('Success')}
  553. # Passed by check_document_access_permission but unused by APIs
  554. notebook = json.loads(request.POST.get('notebook', '{}'))
  555. snippet = json.loads(request.POST.get('snippet', '{}'))
  556. data_format = json.loads(request.POST.get('format', '"hdfs-file"'))
  557. destination = urllib_unquote(json.loads(request.POST.get('destination', '""')))
  558. overwrite = json.loads(request.POST.get('overwrite', 'false'))
  559. is_embedded = json.loads(request.POST.get('is_embedded', 'false'))
  560. start_time = json.loads(request.POST.get('start_time', '-1'))
  561. api = get_api(request, snippet)
  562. if data_format == 'hdfs-file': # Blocking operation, like downloading
  563. if request.fs.isdir(destination):
  564. if notebook.get('name'):
  565. destination += '/%(name)s.csv' % notebook
  566. else:
  567. destination += '/%(type)s-%(id)s.csv' % notebook
  568. if overwrite and request.fs.exists(destination):
  569. request.fs.do_as_user(request.user.username, request.fs.rmtree, destination)
  570. response['watch_url'] = api.export_data_as_hdfs_file(snippet, destination, overwrite)
  571. response['status'] = 0
  572. request.audit = {
  573. 'operation': 'EXPORT',
  574. 'operationText': 'User %s exported to HDFS destination: %s' % (request.user.username, destination),
  575. 'allowed': True
  576. }
  577. elif data_format == 'hive-table':
  578. if is_embedded:
  579. sql, success_url = api.export_data_as_table(notebook, snippet, destination)
  580. task = make_notebook(
  581. name=_('Export %s query to table %s') % (snippet['type'], destination),
  582. description=_('Query %s to %s') % (_get_snippet_name(notebook), success_url),
  583. editor_type=snippet['type'],
  584. statement=sql,
  585. status='ready',
  586. database=snippet['database'],
  587. on_success_url=success_url,
  588. last_executed=start_time,
  589. is_task=True
  590. )
  591. response = task.execute(request)
  592. else:
  593. notebook_id = notebook['id'] or request.GET.get('editor', request.GET.get('notebook'))
  594. response['watch_url'] = reverse('notebook:execute_and_watch') + '?action=save_as_table&notebook=' + str(notebook_id) + '&snippet=0&destination=' + destination
  595. response['status'] = 0
  596. request.audit = {
  597. 'operation': 'EXPORT',
  598. 'operationText': 'User %s exported to Hive table: %s' % (request.user.username, destination),
  599. 'allowed': True
  600. }
  601. elif data_format == 'hdfs-directory':
  602. if destination.lower().startswith("abfs"):
  603. destination = abfspath(destination)
  604. if is_embedded:
  605. sql, success_url = api.export_large_data_to_hdfs(notebook, snippet, destination)
  606. task = make_notebook(
  607. name=_('Export %s query to directory') % snippet['type'],
  608. description=_('Query %s to %s') % (_get_snippet_name(notebook), success_url),
  609. editor_type=snippet['type'],
  610. statement=sql,
  611. status='ready-execute',
  612. database=snippet['database'],
  613. on_success_url=success_url,
  614. last_executed=start_time,
  615. is_task=True
  616. )
  617. response = task.execute(request)
  618. else:
  619. notebook_id = notebook['id'] or request.GET.get('editor', request.GET.get('notebook'))
  620. response['watch_url'] = reverse('notebook:execute_and_watch') + '?action=insert_as_query&notebook=' + str(notebook_id) + '&snippet=0&destination=' + destination
  621. response['status'] = 0
  622. request.audit = {
  623. 'operation': 'EXPORT',
  624. 'operationText': 'User %s exported to HDFS directory: %s' % (request.user.username, destination),
  625. 'allowed': True
  626. }
  627. elif data_format in ('search-index', 'dashboard'):
  628. # Open the result in the Dashboard via a SQL sub-query or the Import wizard (quick vs scalable)
  629. if is_embedded:
  630. notebook_id = notebook['id'] or request.GET.get('editor', request.GET.get('notebook'))
  631. if data_format == 'dashboard':
  632. engine = notebook['type'].replace('query-', '')
  633. response['watch_url'] = reverse('dashboard:browse', kwargs={'name': notebook_id}) + '?source=query&engine=%(engine)s' % {'engine': engine}
  634. response['status'] = 0
  635. else:
  636. sample = get_api(request, snippet).fetch_result(notebook, snippet, rows=4, start_over=True)
  637. for col in sample['meta']:
  638. col['type'] = HiveFormat.FIELD_TYPE_TRANSLATE.get(col['type'], 'string')
  639. response['status'] = 0
  640. response['id'] = notebook_id
  641. response['name'] = _get_snippet_name(notebook)
  642. response['source_type'] = 'query'
  643. response['target_type'] = 'index'
  644. response['target_path'] = destination
  645. response['sample'] = list(sample['data'])
  646. response['columns'] = [
  647. Field(col['name'], col['type']).to_dict() for col in sample['meta']
  648. ]
  649. else:
  650. notebook_id = notebook['id'] or request.GET.get('editor', request.GET.get('notebook'))
  651. response['watch_url'] = reverse('notebook:execute_and_watch') + '?action=index_query&notebook=' + str(notebook_id) + '&snippet=0&destination=' + destination
  652. response['status'] = 0
  653. if response.get('status') != 0:
  654. response['message'] = _('Exporting result failed.')
  655. return JsonResponse(response)
  656. @require_POST
  657. @check_document_access_permission
  658. @api_error_handler
  659. def statement_risk(request):
  660. response = {'status': -1, 'message': ''}
  661. notebook = json.loads(request.POST.get('notebook', '{}'))
  662. snippet = json.loads(request.POST.get('snippet', '{}'))
  663. api = HS2Api(request.user, snippet)
  664. response['query_complexity'] = api.statement_risk(notebook, snippet)
  665. response['status'] = 0
  666. return JsonResponse(response)
  667. @require_POST
  668. @check_document_access_permission
  669. @api_error_handler
  670. def statement_compatibility(request):
  671. response = {'status': -1, 'message': ''}
  672. notebook = json.loads(request.POST.get('notebook', '{}'))
  673. snippet = json.loads(request.POST.get('snippet', '{}'))
  674. source_platform = request.POST.get('sourcePlatform')
  675. target_platform = request.POST.get('targetPlatform')
  676. api = get_api(request, snippet)
  677. response['query_compatibility'] = api.statement_compatibility(notebook, snippet, source_platform=source_platform, target_platform=target_platform)
  678. response['status'] = 0
  679. return JsonResponse(response)
  680. @require_POST
  681. @check_document_access_permission
  682. @api_error_handler
  683. def statement_similarity(request):
  684. response = {'status': -1, 'message': ''}
  685. notebook = json.loads(request.POST.get('notebook', '{}'))
  686. snippet = json.loads(request.POST.get('snippet', '{}'))
  687. source_platform = request.POST.get('sourcePlatform')
  688. api = get_api(request, snippet)
  689. response['statement_similarity'] = api.statement_similarity(notebook, snippet, source_platform=source_platform)
  690. response['status'] = 0
  691. return JsonResponse(response)
  692. @require_POST
  693. @check_document_access_permission
  694. @api_error_handler
  695. def get_external_statement(request):
  696. response = {'status': -1, 'message': ''}
  697. notebook = json.loads(request.POST.get('notebook', '{}'))
  698. snippet = json.loads(request.POST.get('snippet', '{}'))
  699. if snippet.get('statementType') == 'file':
  700. response['statement'] = _get_statement_from_file(request.user, request.fs, snippet)
  701. elif snippet.get('statementType') == 'document':
  702. notebook = Notebook(Document2.objects.get_by_uuid(user=request.user, uuid=snippet['associatedDocumentUuid'], perm_type='read'))
  703. response['statement'] = notebook.get_str()
  704. response['status'] = 0
  705. return JsonResponse(response)
  706. def _get_statement_from_file(user, fs, snippet):
  707. script_path = snippet['statementPath']
  708. if script_path:
  709. script_path = script_path.replace('hdfs://', '')
  710. if fs.do_as_user(user, fs.isfile, script_path):
  711. return fs.do_as_user(user, fs.read, script_path, 0, 16 * 1024 ** 2)
  712. @require_POST
  713. @api_error_handler
  714. def describe(request, database, table=None, column=None):
  715. response = {'status': -1, 'message': ''}
  716. notebook = json.loads(request.POST.get('notebook', '{}'))
  717. source_type = request.POST.get('source_type', '')
  718. snippet = {'type': source_type}
  719. describe = get_api(request, snippet).describe(notebook, snippet, database, table, column=column)
  720. response.update(describe)
  721. return JsonResponse(response)
  722. def _get_snippet(user, notebook, snippet, operation_id):
  723. if operation_id or not snippet:
  724. nb_doc = Document2.objects.get_by_uuid(user=user, uuid=operation_id or notebook['uuid'])
  725. notebook = Notebook(document=nb_doc).get_data()
  726. snippet = notebook['snippets'][0]
  727. return snippet