Browse Source

HUE-2527 [desktop] Truncate more than 255 columns

Erick Tryzelaar 10 years ago
parent
commit
0add6c792d

+ 17 - 4
desktop/core/src/desktop/lib/export_csvxls.py

@@ -28,6 +28,7 @@ from desktop.lib import i18n
 
 
 LOG = logging.getLogger(__name__)
 LOG = logging.getLogger(__name__)
 MAX_XLS_ROWS = 30000
 MAX_XLS_ROWS = 30000
+MAX_XLS_COLS = 255
 
 
 
 
 def nullify(cell):
 def nullify(cell):
@@ -56,20 +57,32 @@ def dataset(headers, data, encoding=None):
 def create_generator(content_generator, format, encoding=None):
 def create_generator(content_generator, format, encoding=None):
   if format == 'csv':
   if format == 'csv':
     show_headers = True
     show_headers = True
-    while True:
-      headers, data = content_generator.next()
+    for headers, data in content_generator:
       yield dataset(show_headers and headers or None, data, encoding).csv
       yield dataset(show_headers and headers or None, data, encoding).csv
       show_headers = False
       show_headers = False
+
   elif format == 'xls':
   elif format == 'xls':
     headers = None
     headers = None
     data = []
     data = []
     for _headers, _data in content_generator:
     for _headers, _data in content_generator:
       # Forced limit on size from tablib
       # Forced limit on size from tablib
       if len(data) > MAX_XLS_ROWS:
       if len(data) > MAX_XLS_ROWS:
-        data = data[:MAX_XLS_ROWS]
         break
         break
+
+      if _headers and len(_headers) > MAX_XLS_COLS:
+        _headers = _headers[:MAX_XLS_COLS]
+
       headers = _headers
       headers = _headers
-      data.extend(_data)
+
+      for row in _data:
+        if len(row) >= MAX_XLS_COLS:
+          row = row[:MAX_XLS_COLS]
+
+        data.append(row)
+
+    if len(data) > MAX_XLS_ROWS:
+      data = data[:MAX_XLS_ROWS]
+
     yield dataset(headers, data, encoding).xls
     yield dataset(headers, data, encoding).xls
     gc.collect()
     gc.collect()
 
 

+ 36 - 6
desktop/core/src/desktop/lib/test_export_csvxls.py

@@ -17,7 +17,7 @@
 
 
 import tablib
 import tablib
 
 
-from desktop.lib.export_csvxls import create_generator, make_response
+from desktop.lib.export_csvxls import MAX_XLS_ROWS, MAX_XLS_COLS, create_generator, make_response
 from nose.tools import assert_equal
 from nose.tools import assert_equal
 
 
 
 
@@ -25,11 +25,11 @@ def content_generator(header, data):
   yield header, data
   yield header, data
 
 
 def test_export_csv():
 def test_export_csv():
-  header = ["x", "y"]
+  headers = ["x", "y"]
   data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ]
   data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ]
 
 
   # Check CSV
   # Check CSV
-  generator = create_generator(content_generator(header, data), "csv")
+  generator = create_generator(content_generator(headers, data), "csv")
   response = make_response(generator, "csv", "foo")
   response = make_response(generator, "csv", "foo")
   assert_equal("application/csv", response["content-type"])
   assert_equal("application/csv", response["content-type"])
   content = ''.join(response.streaming_content)
   content = ''.join(response.streaming_content)
@@ -37,15 +37,45 @@ def test_export_csv():
   assert_equal("attachment; filename=foo.csv", response["content-disposition"])
   assert_equal("attachment; filename=foo.csv", response["content-disposition"])
 
 
 def test_export_xls():
 def test_export_xls():
-  header = ["x", "y"]
+  headers = ["x", "y"]
   data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ]
   data = [ ["1", "2"], ["3", "4"], ["5,6", "7"], [None, None] ]
 
 
-  dataset = tablib.Dataset(headers=header)
+  dataset = tablib.Dataset(headers=headers)
   for row in data:
   for row in data:
     dataset.append([cell is not None and cell or "NULL" for cell in row])
     dataset.append([cell is not None and cell or "NULL" for cell in row])
 
 
   # Check XLS
   # Check XLS
-  generator = create_generator(content_generator(header, data), "xls")
+  generator = create_generator(content_generator(headers, data), "xls")
+  response = make_response(generator, "xls", "foo")
+  assert_equal("application/xls", response["content-type"])
+  content = ''.join(response.streaming_content)
+  assert_equal(dataset.xls, content)
+  assert_equal("attachment; filename=foo.xls", response["content-disposition"])
+
+def test_export_xls_truncate_rows():
+  headers = ["a"]
+  data = [["1"]] * (MAX_XLS_ROWS + 1)
+
+  dataset = tablib.Dataset(headers=headers)
+  dataset.extend(data[:MAX_XLS_ROWS])
+
+  # Check XLS
+  generator = create_generator(content_generator(headers, data), "xls")
+  response = make_response(generator, "xls", "foo")
+  assert_equal("application/xls", response["content-type"])
+  content = ''.join(response.streaming_content)
+  assert_equal(dataset.xls, content)
+  assert_equal("attachment; filename=foo.xls", response["content-disposition"])
+
+def test_export_xls_truncate_cols():
+  headers = ["a"] * (MAX_XLS_COLS + 1)
+  data = [["1"] * (MAX_XLS_COLS + 1)]
+
+  dataset = tablib.Dataset(headers=headers[:MAX_XLS_COLS])
+  dataset.extend([data[0][:MAX_XLS_COLS]])
+
+  # Check XLS
+  generator = create_generator(content_generator(headers, data), "xls")
   response = make_response(generator, "xls", "foo")
   response = make_response(generator, "xls", "foo")
   assert_equal("application/xls", response["content-type"])
   assert_equal("application/xls", response["content-type"])
   content = ''.join(response.streaming_content)
   content = ''.join(response.streaming_content)