Эх сурвалжийг харах

HUE-9208 [editor] Port close statement to v2 and sql task

Romain 5 жил өмнө
parent
commit
214dd41369

+ 7 - 1
desktop/core/src/desktop/js/api/apiHelper.js

@@ -2148,7 +2148,13 @@ class ApiHelper {
    */
   async closeStatement(options) {
     return new Promise(async (resolve, reject) => {
-      simplePost('/notebook/api/close_statement', await options.executable.toContext(), options)
+      simplePost(
+        '/notebook/api/close_statement',
+        {
+          operationId: options.executable.operationId
+        },
+        options
+      )
         .done(resolve)
         .fail(reject);
     });

+ 4 - 1
desktop/libs/notebook/src/notebook/api.py

@@ -665,7 +665,10 @@ def close_statement(request):
 
   notebook = json.loads(request.POST.get('notebook', '{}'))
   snippet = None
-  operation_id = request.POST.get('operationId') or notebook['uuid']
+  operation_id = request.POST.get('operationId')
+
+  if operation_id and not notebook.get('uuid'):
+    notebook['uuid'] = operation_id
 
   snippet = _get_snippet(request.user, notebook, snippet, operation_id)
 

+ 23 - 2
desktop/libs/notebook/src/notebook/task_tests.py

@@ -26,7 +26,7 @@ from desktop.lib.django_test_util import make_logged_in_client
 from useradmin.models import User
 
 from notebook.connectors.sql_alchemy import SqlAlchemyApi
-from notebook.tasks import run_sync_query, download_to_file
+from notebook.tasks import run_sync_query, download_to_file, close_statement
 
 if sys.version_info[0] > 2:
   from unittest.mock import patch, Mock, MagicMock
@@ -63,13 +63,34 @@ class TestRunAsyncQueryTask():
               }.get(key, Mock())
             notebook = MagicMock()
             notebook.__getitem__.side_effect = notebook_dict
-
             snippet = MagicMock()
+
             meta = download_to_file(notebook, snippet)
 
             assert_equal(meta['row_counter'], 2, meta)
 
 
+  def test_close_statement(self):
+    with patch('notebook.tasks._get_request') as _get_request:
+      with patch('notebook.tasks.download_to_file') as download_to_file:
+
+        download_to_file.AsyncResult.return_value = Mock(
+          state=states.SUCCESS
+        )
+
+        def notebook_dict(key):
+          return {
+            'uuid': '1ca47e0d-4708-4709-82c1-a9280e15452b',
+          }.get(key, Mock())
+        notebook = MagicMock()
+        notebook.__getitem__.side_effect = notebook_dict
+        snippet = MagicMock()
+
+        response = close_statement(notebook, snippet)
+
+        assert_equal(response, {'status': 0})
+
+
 class TestRunSyncQueryTask():
 
   def setUp(self):

+ 9 - 2
desktop/libs/notebook/src/notebook/tasks.py

@@ -446,10 +446,13 @@ def cancel(*args, **kwargs):
 
   return {'status': status}
 
+
 def close_statement(*args, **kwargs):
   notebook = args[0]
   snippet = args[1]
-  result = download_to_file.AsyncResult(notebook['uuid'])
+  task_id = _get_query_key(notebook, snippet)
+
+  result = download_to_file.AsyncResult(task_id)
   state = result.state
   status = 0
 
@@ -466,12 +469,16 @@ def close_statement(*args, **kwargs):
     close_statement_async.apply_async(args=args, kwargs=kwargs, task_id=_close_statement_async_id(notebook, snippet))
 
   result.forget()
+
   _cleanup(notebook, snippet)
 
   return {'status': status}
 
+
 def _cleanup(notebook, snippet):
-  storage.delete(_result_key(notebook, snippet)) # TODO: abstract storage + caches
+  task_id = _get_query_key(notebook, snippet)
+
+  storage.delete(_result_key(task_id))  # TODO: abstract storage + caches
   storage.delete(_log_key(notebook, snippet))
   caches[CACHES_CELERY_KEY].delete(_fetch_progress_key(notebook, snippet))