Browse Source

HUE-5165 [impala] Get result count for all select * queries

Jenny Kim 9 years ago
parent
commit
d89587b

+ 10 - 3
desktop/libs/notebook/src/notebook/connectors/hiveserver2.py

@@ -21,6 +21,7 @@ import logging
 import re
 import StringIO
 import struct
+import time
 
 from django.core.urlresolvers import reverse
 from django.utils.translation import ugettext as _
@@ -734,9 +735,15 @@ class HS2Api(Api):
     server_url = '%s://%s' % (protocol, self._get_impala_server_url(session))
     if query_id:
       LOG.info("Attempting to get Impala query profile at server_url %s for query ID: %s" % (server_url, query_id))
-      fragment = self._get_impala_query_profile(server_url, query_id=query_id)
-      total_records_re = "Coordinator Fragment F\d\d.+?RowsReturned: \d+(?:.\d+[KMB])? \((?P<total_records>\d+)\).*?Averaged Fragment F\d\d"
-      total_records_match = re.search(total_records_re, fragment, re.MULTILINE | re.DOTALL)
+      retries = 0
+      max_retries = 10
+      total_records_match = None
+      while not total_records_match and retries < max_retries:
+        time.sleep(1.0)
+        fragment = self._get_impala_query_profile(server_url, query_id=query_id)
+        total_records_re = "Coordinator Fragment F\d\d.+?RowsReturned: \d+(?:.\d+[KMB])? \((?P<total_records>\d+)\).*?(Averaged Fragment F\d\d)"
+        total_records_match = re.search(total_records_re, fragment, re.MULTILINE | re.DOTALL)
+        retries += 1
     if total_records_match:
       total_records = int(total_records_match.group('total_records'))
 

+ 18 - 0
desktop/libs/notebook/src/notebook/connectors/tests/tests_hiveserver2.py

@@ -756,6 +756,24 @@ class TestHiveserver2ApiWithHadoop(BeeswaxSampleProvider):
       assert_true('size' in data['result'])
       assert_equal(23, data['result']['rows'])
       assert_equal(None, data['result']['size'])
+
+      # Assert that selecting all from partitioned table works
+      statement = "SELECT * FROM web_logs;"
+      doc = self.create_query_document(owner=self.user, query_type='impala', statement=statement)
+      notebook = Notebook(document=doc)
+      snippet = self.execute_and_wait(doc, snippet_idx=0, timeout=60.0, wait=2.0)
+
+      self.client.post(reverse('notebook:fetch_result_data'),
+                       {'notebook': notebook.get_json(), 'snippet': json.dumps(snippet), 'rows': 100, 'startOver': 'false'})
+
+      response = self.client.post(reverse('notebook:fetch_result_size'),
+                                 {'notebook': notebook.get_json(), 'snippet': json.dumps(snippet)})
+
+      data = json.loads(response.content)
+      assert_equal(0, data['status'], data)
+      assert_true('result' in data)
+      assert_true('rows' in data['result'])
+      assert_equal(1000, data['result']['rows'])
     finally:
       self.api.close_session(session)