Sfoglia il codice sorgente

HUE-8747 [editor] Fix result count off by one.

jdesjean 6 anni fa
parent
commit
a2203ac610

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

@@ -130,7 +130,6 @@ class DataAdapter:
 
   def next(self):
     results = self.db.fetch(start_over=self.start_over, rows=self.fetch_size)
-
     if self.first_fetched:
       self.first_fetched = False
       self.start_over = False
@@ -154,18 +153,17 @@ class DataAdapter:
       data = []
 
       for row in results['data']:
-        self.row_counter += 1
-        if self.limit_bytes:
-          self.bytes_counter += self._getsizeofascii(row)
-
-        if self.limit_rows and self.row_counter > self.max_rows:
+        num_bytes = self._getsizeofascii(row)
+        if self.limit_rows and self.row_counter + 1 > self.max_rows:
           LOG.warn('The query results exceeded the maximum row limit of %d and has been truncated to first %d rows.' % (self.max_rows, self.row_counter))
           self.is_truncated = True
           break
-        if self.limit_bytes and self.bytes_counter > self.max_bytes:
+        if self.limit_bytes and self.bytes_counter + num_bytes > self.max_bytes:
           LOG.warn('The query results exceeded the maximum bytes limit of %d and has been truncated to first %d rows.' % (self.max_bytes, self.row_counter))
           self.is_truncated = True
           break
+        self.row_counter += 1
+        self.bytes_counter += num_bytes
         data.append(row)
 
       return self.headers, data

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

@@ -95,7 +95,7 @@ def download_to_file(notebook, snippet, file_format='csv', max_rows=-1, **kwargs
   f_log, path_log = tempfile.mkstemp()
   try:
     #TODO: We need to move this metadata somewhere else, it gets erased on exception and we can no longer cleanup the files.
-    meta = {'row_counter': 0, 'file_path': path, 'handle': {}, 'log_path': path_log, 'status': 'running', 'truncated': False} #TODO: Truncated
+    meta = {'row_counter': 0, 'file_path': path, 'handle': {}, 'log_path': path_log, 'status': 'running', 'truncated': False}
 
     result_wrapper = ResultWrapper(api, notebook, snippet, ResultWrapperCallback(notebook['uuid'], meta, f_log))
     content_generator = data_export.DataAdapter(result_wrapper, max_rows=max_rows, store_data_type_in_header=True) #TODO: Move PREFETCH_RESULT_COUNT to front end
@@ -104,6 +104,7 @@ def download_to_file(notebook, snippet, file_format='csv', max_rows=-1, **kwargs
     for chunk in response:
       os.write(f, chunk)
       meta['row_counter'] = content_generator.row_counter
+      meta['truncated'] = content_generator.is_truncated
       download_to_file.update_state(task_id=notebook['uuid'], state='AVAILABLE', meta=meta)
 
   finally: