Browse Source

HUE-5679 [impala] Do not truncate the last part of rows in result downloads

Romain Rigaux 8 years ago
parent
commit
44cbbb30d3
2 changed files with 49 additions and 3 deletions
  1. 5 3
      apps/beeswax/src/beeswax/data_export.py
  2. 44 0
      apps/impala/src/impala/tests.py

+ 5 - 3
apps/beeswax/src/beeswax/data_export.py

@@ -93,6 +93,7 @@ class HS2DataAdapter:
     self.num_cols = None
     self.num_cols = None
     self.row_counter = 1
     self.row_counter = 1
     self.is_truncated = False
     self.is_truncated = False
+    self.has_more = True
 
 
   def __iter__(self):
   def __iter__(self):
     return self
     return self
@@ -101,6 +102,8 @@ class HS2DataAdapter:
     results = self.db.fetch(self.handle, start_over=self.start_over, rows=self.fetch_size)
     results = self.db.fetch(self.handle, start_over=self.start_over, rows=self.fetch_size)
 
 
     if self.first_fetched:
     if self.first_fetched:
+      self.first_fetched = False
+      self.start_over = False
       self.headers = results.cols()
       self.headers = results.cols()
       self.num_cols = len(self.headers)
       self.num_cols = len(self.headers)
 
 
@@ -109,9 +112,8 @@ class HS2DataAdapter:
         LOG.warn('The query results contain %d columns and may take long time to download, reducing fetch size to 100.' % self.num_cols)
         LOG.warn('The query results contain %d columns and may take long time to download, reducing fetch size to 100.' % self.num_cols)
         self.fetch_size = 100
         self.fetch_size = 100
 
 
-    if not self.is_truncated and (self.first_fetched or results.has_more):
-      self.first_fetched = False
-      self.start_over = False
+    if self.has_more and not self.is_truncated:
+      self.has_more = results.has_more
       data = []
       data = []
 
 
       for row in results.rows():
       for row in results.rows():

+ 44 - 0
apps/impala/src/impala/tests.py

@@ -31,7 +31,10 @@ from desktop.lib.test_utils import add_to_group
 from desktop.models import Document
 from desktop.models import Document
 from hadoop.pseudo_hdfs4 import get_db_prefix, is_live_cluster
 from hadoop.pseudo_hdfs4 import get_db_prefix, is_live_cluster
 
 
+from beeswax import data_export
 from beeswax.design import hql_query
 from beeswax.design import hql_query
+
+from beeswax.data_export import download
 from beeswax.models import SavedQuery, QueryHistory
 from beeswax.models import SavedQuery, QueryHistory
 from beeswax.server import dbms
 from beeswax.server import dbms
 from beeswax.test_base import get_query_server_config, wait_for_query_to_finish, fetch_query_result_data
 from beeswax.test_base import get_query_server_config, wait_for_query_to_finish, fetch_query_result_data
@@ -209,6 +212,47 @@ class TestImpalaIntegration:
     content = json.loads(resp.content)
     content = json.loads(resp.content)
     assert_equal(0, content['status'])
     assert_equal(0, content['status'])
 
 
+
+  def test_data_download(self):
+    hql = 'SELECT * FROM tweets %(limit)s'
+
+    FETCH_SIZE = data_export.FETCH_SIZE
+    data_export.FETCH_SIZE = 2 # Decrease fetch size to validate last fetch logic
+
+    try:
+      query = hql_query(hql % {'limit': ''})
+
+      handle = self.db.execute_and_wait(query)
+      # Get the result in csv. Should have 5 + 1 header row.
+      csv_resp = download(handle, 'csv', self.db)
+      csv_content = ''.join(csv_resp.streaming_content)
+      assert_equal(len(csv_content.strip().split('\n')), 5 + 1)
+
+
+      query = hql_query(hql % {'limit': 'LIMIT 0'})
+
+      handle = self.db.execute_and_wait(query)
+      csv_resp = download(handle, 'csv', self.db)
+      csv_content = ''.join(csv_resp.streaming_content)
+      assert_equal(len(csv_content.strip().split('\n')), 1)
+
+      query = hql_query(hql % {'limit': 'LIMIT 1'})
+
+      handle = self.db.execute_and_wait(query)
+      csv_resp = download(handle, 'csv', self.db)
+      csv_content = ''.join(csv_resp.streaming_content)
+      assert_equal(len(csv_content.strip().split('\n')), 1 + 1)
+
+      query = hql_query(hql % {'limit': 'LIMIT 2'})
+
+      handle = self.db.execute_and_wait(query)
+      csv_resp = download(handle, 'csv', self.db)
+      csv_content = ''.join(csv_resp.streaming_content)
+      assert_equal(len(csv_content.strip().split('\n')), 1 + 2)
+    finally:
+      data_export.FETCH_SIZE = FETCH_SIZE
+
+
   def test_explain(self):
   def test_explain(self):
     QUERY = """
     QUERY = """
       SELECT * FROM tweets ORDER BY row_num;
       SELECT * FROM tweets ORDER BY row_num;