فهرست منبع

HUE-9162 [sqlalchemy] Avoid API error on data sampling popup

Should always return a 'row' item list, even if empty.
Romain 5 سال پیش
والد
کامیت
3bb43032fc

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

@@ -220,10 +220,7 @@ def execute(request, engine=None):
 
     response = _execute_notebook(request, notebook, snippet)
 
-    span.set_tag(
-      'query-id',
-      response['handle']['guid'] if response.get('handle') and response['handle'].get('guid') else None
-    )
+    span.set_tag('query-id', response.get('handle', {}).get('guid'))
 
   return JsonResponse(response)
 

+ 2 - 4
desktop/libs/notebook/src/notebook/connectors/sql_alchemy.py

@@ -303,11 +303,9 @@ class SqlAlchemyApi(Api):
     response = {'status': -1, 'result': {}}
 
     metadata, sample_data = assist.get_sample_data(database, table, column=column, operation=operation)
-    has_result_set = sample_data is not None
 
-    if sample_data:
-      response['status'] = 0
-      response['rows'] = escape_rows(sample_data)
+    response['status'] = 0
+    response['rows'] = escape_rows(sample_data)
 
     if table and operation != 'hello':
       columns = assist.get_columns(database, table)

+ 28 - 22
desktop/libs/notebook/src/notebook/connectors/sql_alchemy_tests.py

@@ -28,6 +28,7 @@ from useradmin.models import User
 
 from notebook.connectors.sql_alchemy import SqlAlchemyApi
 
+
 if sys.version_info[0] > 2:
   from unittest.mock import patch, Mock, MagicMock
 else:
@@ -43,6 +44,11 @@ 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"))
+    self.interpreter = {
+      'options': {
+        'url': 'mysql://hue:localhost@hue:3306/hue'
+      },
+    }
 
 
   def test_column_backticks_escaping(self):
@@ -64,7 +70,7 @@ class TestApi(object):
   def test_create_athena_engine(self):
     interpreter = {
       'options': {
-        'url': 'awsathena+rest://XXXXXXXXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX@athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://gethue-athena/scratch'
+        'url': 'awsathena+rest://XXXXXXXXXXXXXXX:XXXXXXXXXXXXXXXXXXX@athena.us-west-2.amazonaws.com:443/default?s3_staging_dir=s3://gethue-athena/scratch'
       }
     }
 
@@ -73,12 +79,6 @@ class TestApi(object):
 
 
   def test_fetch_result_empty(self):
-    interpreter = {
-      'options': {
-        'url': 'mysql://hue:localhost@hue:3306/hue'
-      },
-    }
-
     notebook = Mock()
     snippet = {'result': {'handle': {'guid': 'guid-1'}}}
     rows = 10
@@ -97,7 +97,7 @@ class TestApi(object):
         }
       )
 
-      data = SqlAlchemyApi(self.user, interpreter).fetch_result(notebook, snippet, rows, start_over)
+      data = SqlAlchemyApi(self.user, self.interpreter).fetch_result(notebook, snippet, rows, start_over)
 
       assert_false(data['has_more'])
       assert_not_equal(data['has_more'], [])
@@ -108,12 +108,6 @@ class TestApi(object):
 
 
   def test_fetch_result_rows(self):
-    interpreter = {
-      'options': {
-        'url': 'mysql://hue:localhost@hue:3306/hue'
-      }
-    }
-
     notebook = Mock()
     snippet = {'result': {'handle': {'guid': 'guid-1'}}}
     rows = 10
@@ -132,7 +126,7 @@ class TestApi(object):
         }
       )
 
-      data = SqlAlchemyApi(self.user, interpreter).fetch_result(notebook, snippet, rows, start_over)
+      data = SqlAlchemyApi(self.user, self.interpreter).fetch_result(notebook, snippet, rows, start_over)
 
       assert_false(data['has_more'])
       assert_not_equal(data['has_more'], [])
@@ -141,25 +135,37 @@ 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)
+      response = SqlAlchemyApi(self.user, self.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)
+      response = SqlAlchemyApi(self.user, self.interpreter).check_status(notebook, snippet)
       assert_equal(response['status'], 'available')
 
 
+  def test_get_sample_data(self):
+    snippet = Mock()
+
+    with patch('notebook.connectors.sql_alchemy.Assist.get_sample_data') as get_sample_data:
+      with patch('notebook.connectors.sql_alchemy.inspect') as inspect:
+        get_sample_data.return_value = (['col1'], [[1], [2]])
+
+        response = SqlAlchemyApi(self.user, self.interpreter).get_sample_data(snippet)
+
+        assert_equal(response['rows'], [[1], [2]])
+        assert_equal(
+          response['full_headers'],
+          [{'name': 'col1', 'type': 'STRING_TYPE', 'comment': ''}]
+        )
+
+
 class TestAutocomplete(object):
 
   def setUp(self):