Răsfoiți Sursa

HUE-9162 [sqlalchemy] Do not fetch result on query without a resultset

Romain 5 ani în urmă
părinte
comite
25e7b54ef1

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

@@ -242,7 +242,7 @@ def check_status(request):
     span.set_tag('user-id', request.user.username)
     span.set_tag(
       'query-id',
-      snippet['result']['handle']['guid'] if snippet['result'].get('handle') and snippet['result']['handle'].get('guid') else None
+      snippet.get('result', {}).get('handle', {}).get('guid')
     )
 
     response = _check_status(request, notebook=notebook, snippet=snippet, operation_id=operation_id)

+ 8 - 3
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -169,10 +169,15 @@ class SqlAlchemyApi(Api):
     guid = snippet['result']['handle']['guid']
     connection = CONNECTION_CACHE.get(guid)
 
+    response = {'status': 'canceled'}
+
     if connection:
-      return {'status': 'available'}
-    else:
-      return {'status': 'canceled'}
+      if snippet['result']['handle']['has_result_set']:
+        response['status'] = 'available'
+      else:
+        response['status'] = 'success'
+
+    return response
 
   @query_error_handler
   def fetch_result(self, notebook, snippet, rows, start_over):

+ 19 - 3
desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -24,7 +24,6 @@ from nose.tools import assert_equal, assert_not_equal, assert_true, assert_false
 
 from desktop.auth.backend import rewrite_user
 from desktop.lib.django_test_util import make_logged_in_client
-from desktop.lib.test_utils import add_to_group, grant_access
 from useradmin.models import User
 
 from notebook.connectors.sql_alchemy import SqlAlchemyApi
@@ -34,6 +33,7 @@ if sys.version_info[0] > 2:
 else:
   from mock import patch, Mock, MagicMock
 
+
 LOG = logging.getLogger(__name__)
 
 
@@ -43,7 +43,6 @@ class TestApi(object):
     self.client = make_logged_in_client(username="test", groupname="default", recreate=True, is_superuser=False)
 
     self.user = rewrite_user(User.objects.get(username="test"))
-    grant_access("test", "default", "notebook")
 
 
   def test_column_backticks_escaping(self):
@@ -142,6 +141,24 @@ class TestApi(object):
       assert_equal(data['data'], [['row1'], ['row2']])
       assert_equal(data['meta'](), [{'type': 'BIGINT_TYPE'}])
 
+  def test_check_status(self):
+    interpreter = {
+      'options': {
+        'url': 'mysql://hue:localhost@hue:3306/hue'
+      },
+    }
+    notebook = Mock()
+
+    with patch('notebook.connectors.sql_alchemy.CONNECTION_CACHE') as CONNECTION_CACHE:
+
+      snippet = {'result': {'handle': {'guid': 'guid-1', 'has_result_set': False}}}
+      response = SqlAlchemyApi(self.user, interpreter).check_status(notebook, snippet)
+      assert_equal(response['status'], 'success')
+
+      snippet = {'result': {'handle': {'guid': 'guid-1', 'has_result_set': True}}}
+      response = SqlAlchemyApi(self.user, interpreter).check_status(notebook, snippet)
+      assert_equal(response['status'], 'available')
+
 
 class TestAutocomplete(object):
 
@@ -149,7 +166,6 @@ class TestAutocomplete(object):
     self.client = make_logged_in_client(username="test", groupname="default", recreate=True, is_superuser=False)
 
     self.user = rewrite_user(User.objects.get(username="test"))
-    grant_access("test", "default", "notebook")
 
 
   def test_empty_database_names(self):