Эх сурвалжийг харах

[core] Stream CSV download and limit XLS download

CSV download can be streamed because that format can be contatenated easily.
Tablib XLS has a max size of 65000 rows, so it must be limited.
Also added a "starting yield" so that downloads start early rather than later.
Abraham Elmahrek 11 жил өмнө
parent
commit
e436a8a11f

+ 16 - 19
apps/beeswax/src/beeswax/data_export.py

@@ -43,8 +43,9 @@ def download(handle, format, db):
     LOG.error('Unknown download format "%s"' % (format,))
     return
 
-  data, has_more = HS2DataAdapter(handle, db, conf.DOWNLOAD_ROW_LIMIT.get())
-  return export_csvxls.make_response(data[0], data[1:], format, 'query_result')
+  content_generator = HS2DataAdapter(handle, db, conf.DOWNLOAD_ROW_LIMIT.get())
+  generator = export_csvxls.create_generator(content_generator, format)
+  return export_csvxls.make_response(generator, format, 'query_result')
 
 
 def upload(path, handle, user, db, fs):
@@ -53,28 +54,20 @@ def upload(path, handle, user, db, fs):
 
   Retrieve the query result in the format specified and upload to hdfs.
   """
-  has_more = True
-  start_over = True
-
   if fs.do_as_user(user.username, fs.exists, path):
     raise Exception(_("%s already exists.") % path)
   else:
     fs.do_as_user(user.username, fs.create, path)
 
-  while has_more:
-    data, has_more = HS2DataAdapter(handle, db, conf.DOWNLOAD_ROW_LIMIT.get(), start_over=start_over)
-    dataset = export_csvxls.dataset(None, data[1:])
+  content_generator = HS2DataAdapter(handle, db, -1, start_over=True)
+  for header, data in content_generator:
+    dataset = export_csvxls.dataset(None, data)
     fs.do_as_user(user.username, fs.append, path, dataset.csv)
 
-    if start_over:
-      start_over = False
-
 
 def HS2DataAdapter(handle, db, max_rows=0, start_over=True):
   """
-  HS2DataAdapter(query_model, db) -> 2D array of data.
-
-  First line should be the headers.
+  HS2DataAdapter(query_model, db) -> headers, 2D array of data.
   """
 
   results = db.fetch(handle, start_over=start_over, rows=FETCH_SIZE)
@@ -83,20 +76,24 @@ def HS2DataAdapter(handle, db, max_rows=0, start_over=True):
     time.sleep(_DATA_WAIT_SLEEP)
     results = db.fetch(handle, start_over=start_over, rows=FETCH_SIZE)
 
-  data = [results.cols()]
+  headers = results.cols()
 
+  num_rows_seen = 0
+  limit_rows = max_rows > -1
   while results is not None:
+    data = []
     for row in results.rows():
-      if max_rows > -1 and len(data) > max_rows:
+      num_rows_seen += 1
+      if limit_rows and num_rows_seen > max_rows:
         break
       data.append(row)
 
-    if max_rows > -1 and len(data) > max_rows:
+    yield headers, data
+
+    if limit_rows and num_rows_seen > max_rows:
       break
 
     if results.has_more:
       results = db.fetch(handle, start_over=False, rows=FETCH_SIZE)
     else:
       results = None
-
-  return data, results.has_more if results else False

+ 6 - 7
apps/search/src/search/data_export.py

@@ -38,20 +38,19 @@ def download(results, format):
     LOG.error('Unknown download format "%s"' % format)
     return
 
-  data = SearchDataAdapter(results, format)
-  return export_csvxls.make_response(data[0], data[1:], format, 'query_result')
+  content_generator = SearchDataAdapter(results, format)
+  generator = export_csvxls.create_generator(content_generator, format)
+  return export_csvxls.make_response(generator, format, 'query_result')
 
 
 def SearchDataAdapter(results, format):
   """
-  SearchDataAdapter(results, format, db) -> 2D array of data.
-
-  First line should be the headers.
+  SearchDataAdapter(results, format, db) -> headers, 2D array of data.
   """
   if results and results['response'] and results['response']['docs']:
     search_data = results['response']['docs']
     order = search_data[0].keys()
-    rows = [order]
+    rows = []
 
     for data in search_data:
       row = []
@@ -66,4 +65,4 @@ def SearchDataAdapter(results, format):
   else:
     rows = [[]]
 
-  return rows
+  yield order, rows

Файлын зөрүү хэтэрхий том тул дарагдсан байна
+ 0 - 0
apps/search/src/search/tests.py


+ 27 - 11
desktop/core/src/desktop/lib/export_csvxls.py

@@ -17,6 +17,7 @@
 """
 Common library to export either CSV or XLS.
 """
+import gc
 import logging
 import tablib
 
@@ -24,11 +25,10 @@ from django.http import HttpResponse
 from django.utils.encoding import smart_str
 from desktop.lib import i18n
 
+
 LOG = logging.getLogger(__name__)
-XLS_SIZE_LIMIT = 1024 * 1024 * 1024      # 1GB
+MAX_XLS_ROWS = 30000
 
-class TooBigToDownloadException(Exception):
-  pass
 
 def nullify(cell):
   return cell if cell is not None else "NULL"
@@ -52,9 +52,30 @@ def dataset(headers, data, encoding=None):
 
   return dataset
 
-def make_response(headers, data, format, name, encoding=None):
+
+def create_generator(content_generator, format, encoding=None):
+  if format == 'csv':
+    show_headers = True
+    while True:
+      headers, data = content_generator.next()
+      yield dataset(show_headers and headers or None, data, encoding).csv
+      show_headers = False
+  elif format == 'xls':
+    headers = None
+    data = []
+    for _headers, _data in content_generator:
+      # Forced limit on size from tablib
+      if len(data) > MAX_XLS_ROWS:
+        data = data[:MAX_XLS_ROWS]
+        break
+      headers = _headers
+      data.extend(_data)
+    yield dataset(headers, data, encoding).xls
+    gc.collect()
+
+
+def make_response(generator, format, name, encoding=None):
   """
-  @param headers List of strings to form the header
   @param data An iterator of rows, where every row is a list of strings
   @param format Either "csv" or "xls"
   @param name Base name for output file
@@ -67,11 +88,6 @@ def make_response(headers, data, format, name, encoding=None):
   else:
     raise Exception("Unknown format: %s" % format)
 
-  formatted_data = getattr(dataset(headers, data, encoding), format)
-
-  if len(formatted_data) > XLS_SIZE_LIMIT:
-    raise TooBigToDownloadException()
-
-  resp = HttpResponse(formatted_data, mimetype=mimetype)
+  resp = HttpResponse(generator, mimetype=mimetype)
   resp['Content-Disposition'] = 'attachment; filename=%s.%s' % (name, format)
   return resp

+ 9 - 3
desktop/core/src/desktop/lib/test_export_csvxls.py

@@ -17,15 +17,20 @@
 
 import tablib
 
-from desktop.lib.export_csvxls import make_response
+from desktop.lib.export_csvxls import create_generator, make_response
 from nose.tools import assert_equal
 
+
+def content_generator(header, data):
+  yield header, data
+
 def test_export_csv():
   header = ["x", "y"]
   data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ]
 
   # Check CSV
-  response = make_response(header, data, "csv", "foo")
+  generator = create_generator(content_generator(header, data), "csv")
+  response = make_response(generator, "csv", "foo")
   assert_equal("application/csv", response["content-type"])
   assert_equal('x,y\r\n1,2\r\n3,4\r\n"5,6",7\r\nNULL,NULL\r\n', response.content)
   assert_equal("attachment; filename=foo.csv", response["content-disposition"])
@@ -39,7 +44,8 @@ def test_export_xls():
     dataset.append([cell is not None and cell or "NULL" for cell in row])
 
   # Check XLS
-  response = make_response(header, data, "xls", "foo")
+  generator = create_generator(content_generator(header, data), "xls")
+  response = make_response(generator, "xls", "foo")
   assert_equal("application/xls", response["content-type"])
   assert_equal(dataset.xls, response.content)
   assert_equal("attachment; filename=foo.xls", response["content-disposition"])

Энэ ялгаанд хэт олон файл өөрчлөгдсөн тул зарим файлыг харуулаагүй болно