Explorar el Código

GH-1424 [api] Support Hive query kill action in jobbrowser

sreenaths hace 5 años
padre
commit
e2f186b664

+ 1 - 1
apps/jobbrowser/src/jobbrowser/api2.py

@@ -119,7 +119,7 @@ def action(request, interface=None, action=None):
 
   response['operation'] = operation
   response.update(
-      get_api(request.user, interface, cluster=cluster).action(app_ids, operation)
+      get_api(request.user, interface, cluster=cluster).action(app_ids, operation, request)
   )
 
   return JsonResponse(response)

+ 33 - 13
apps/jobbrowser/src/jobbrowser/apis/hive_query_api.py

@@ -17,7 +17,7 @@
 from builtins import filter
 
 import logging
-import re
+from logging import exception
 
 from datetime import datetime
 from django.utils.translation import ugettext as _
@@ -27,16 +27,17 @@ from desktop.lib.python_util import current_ms_from_utc
 from desktop.lib.rest.http_client import HttpClient
 from desktop.lib.rest.resource import Resource
 
-from notebook.models import _get_notebook_api
+from notebook.models import _get_notebook_api, make_notebook
+
+from beeswax.models import QueryHistory
 
 from jobbrowser.apis.base_api import Api
-from jobbrowser.models import HiveQuery
 from jobbrowser.conf import QUERY_STORE
+from jobbrowser.models import HiveQuery
 
 
 LOG = logging.getLogger(__name__)
 
-
 class HiveQueryApi(Api):
   HEADERS = {'X-Requested-By': 'das'}
 
@@ -114,21 +115,40 @@ class HiveQueryApi(Api):
 
     return app
 
-  def action(self, appid, action):
-    message = {'message': '', 'status': 0}
+  def action(self, query_ids, action, request=None):
+    message = {'actions': {}, 'status': 0}
 
     if action.get('action') == 'kill':
-      for queryid in appid:
-        notebook = {}
-        snippet = {'result': {'handle': {'secret': queryid, 'guid': queryid, 'has_result_set': False}}}
-        connector_id = 'hive'
+      for query_id in query_ids:
+        action_details = {}
+
+        try:
+          self.kill_query(query_id, request)
+          action_details['status'] = 0
+          action_details['message'] = _('kill action performed')
+        except Exception as ex:
+          LOG.error(ex)
+          message['status'] = -1
+          action_details['status'] = -1
+          action_details['message'] = _('kill action failed : %s' % str(ex))
 
-        response = _get_notebook_api(self.user, connector_id).cancel(notebook, snippet)
-        message['status'] = response['status'] if response['status'] != 0 else message['status']
-        message['message'] = _('kill action performed')
+        message['actions'][query_id] = action_details;
 
     return message
 
+  def kill_query(self, query_id, request):
+    kill_sql = 'KILL QUERY "%s";' % query_id
+    job = make_notebook(
+        name=_('Kill query %s') % query_id,
+        editor_type='hive',
+        statement=kill_sql,
+        status='ready',
+        on_success_url='assist.db.refresh',
+        is_task=False,
+    )
+
+    job.execute_and_wait(request)
+
   def logs(self, appid, app_type, log_name=None, is_embeddable=False):
     return {'logs': ''}
 

+ 13 - 10
apps/jobbrowser/src/jobbrowser/apis/hive_query_api_tests.py

@@ -52,24 +52,27 @@ class TestHiveQueryApiNotebook():
 
 
   def test_kill_query(self):
-    with patch('jobbrowser.apis.hive_query_api._get_notebook_api') as _get_notebook_api:
-      cancel_call = Mock(return_value={'status': 0})
-      _get_notebook_api.return_value = Mock(cancel=cancel_call)
+    with patch('jobbrowser.apis.hive_query_api.make_notebook') as make_notebook:
+      execute_and_wait = Mock()
+      make_notebook.return_value = Mock(execute_and_wait=execute_and_wait)
 
-      appid = 'd94d2fb4815a05c4:b1ccec1500000000'
+      query_id = 'hive_20201124114044_bd1b8d39-f18f-4d89-ae1b-7a35e7950579'
       data = {
         'operation': json.dumps({'action': 'kill'}),
         'interface': json.dumps('queries-hive'),
-        'app_ids': json.dumps([appid])
+        'app_ids': json.dumps([query_id])
       }
       response = self.client.post("/jobbrowser/api/job/action/queries-hive/kill", data)
       response_data = json.loads(response.content)
 
-      notebook = {}
-      snippet = {'result': {'handle': {'has_result_set': False, 'secret': appid, 'guid': appid}}}
-
-      _get_notebook_api.assert_called_once_with(self.user, 'hive')
-      cancel_call.assert_called_once_with(notebook, snippet)
+      make_notebook.assert_called_once_with(
+        name='Kill query %s' % query_id,
+        editor_type='hive',
+        statement='KILL QUERY "%s";' % query_id,
+        status='ready',
+        on_success_url='assist.db.refresh',
+        is_task=False,
+      )
 
       assert_equal(0, response_data['status'])