Merge pull request #10 from kubrasadeb/feat/download-zip-file

Download data as zipped file if contains more than 1M rows

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "django-report-builder"
-version = "6.5.7"
+version = "6.6.0"
 description = "Query and Report builder for Django ORM"
 authors = ["David Burke <david@burkesoftware.com>"]
 packages = [

report_builder/mixins.py

@@ -23,6 +23,8 @@
     get_model_from_path_string,
     get_custom_fields_from_model,
 )
+from datetime import datetime
+
 
 DisplayField = namedtuple(
     "DisplayField",
@@ -33,7 +35,7 @@
 def generate_filename(title, ends_with):
     title = title.split('.')[0]
     title.replace(' ', '_')
-    title += ('_' + datetime.datetime.now().strftime("%m%d_%H%M"))
+    title += ('_' + datetime.now().strftime("%m%d_%H%M"))
     if not title.endswith(ends_with):
         title += ends_with
     return title
@@ -124,22 +126,25 @@ def list_to_workbook(self, data, title='report', header=None, widths=None):
             self.build_sheet(data, ws, header=header, widths=widths)
         return wb
 
-    def list_to_xlsx_file(self, data, title='report', header=None, widths=None):
-        """ Make 2D list into a xlsx response for download
-        data can be a 2d array or a dict of 2d arrays
-        like {'sheet_1': [['A1', 'B1']]}
-        returns a StringIO file
-        """
-        wb = self.list_to_workbook(data, title, header, widths)
-        if not title.endswith('.xlsx'):
-            title += '.xlsx'
-        with NamedTemporaryFile() as tmp:
-            wb.save(tmp.name)
-            tmp.seek(0)
-            stream = tmp.read()
-        myfile = BytesIO()
-        myfile.write(stream)
-        return myfile
+    def list_to_xlsx_file(self, data, title, header=None, widths=None):
+        """Make a list into an XLSX file."""
+        wb = Workbook()
+        title = re.sub(r'\W+', '', title)[:30]
+        ws = wb.active
+        if header:
+            ws.append(header)
+        for row in data:
+            cleaned_row = []
+            for value in row:
+                if isinstance(value, datetime):
+                    value = value.replace(tzinfo=None)
+                cleaned_row.append(value)
+            ws.append(cleaned_row)
+
+        file_buffer = BytesIO()
+        wb.save(file_buffer)
+        file_buffer.seek(0)
+        return file_buffer
 
     def list_to_csv_file(self, data, title='report', header=None, widths=None):
         """ Make a list into a csv response for download.
@@ -413,7 +418,7 @@ def increment_total(display_field_key, val):
             defaults = {
                 None: str,
                 datetime.date: lambda: datetime.date(datetime.MINYEAR, 1, 1),
-                datetime.datetime: lambda: datetime.datetime(datetime.MINYEAR, 1, 1),
+                datetime: lambda: datetime(datetime.MINYEAR, 1, 1),
             }
 
             # Order sort fields in reverse order so that ascending, descending

report_builder/models.py

@@ -26,6 +26,9 @@
     increment_total,
     formatter
 )
+import zipfile
+from io import BytesIO
+
 
 AUTH_USER_MODEL = getattr(settings, 'AUTH_USER_MODEL', 'auth.User')
 
@@ -439,27 +442,62 @@ def email_report(self, user=None, email=None):
         report_url = self.report_file.url
         return email_report(report_url, user=user, email=email)
 
-    def async_report_save(self, objects_list,
-                          title, header, widths, user=None, file_type="xlsx", email_to: str = None):
+    def async_report_save(self, chunks, title, header, widths, user=None, file_type=None, email_to: str = None):
+        data_export = DataExportMixin()
+        if file_type not in ["csv", "xlsx"]:
+            raise ValueError("file_type must be 'csv' or 'xlsx'")
         data_export = DataExportMixin()
-        if file_type == 'csv':
-            csv_file = data_export.list_to_csv_file(objects_list, title,
-                                                    header, widths)
-            title = generate_filename(title, '.csv')
-            self.report_file.save(title, ContentFile(csv_file.getvalue().encode()))
+
+        if len(chunks) == 1:
+            single_chunk = chunks[0]
+            if file_type == 'csv':
+                csv_file = data_export.list_to_csv_file(
+                    single_chunk, title, header, widths
+                )
+                file_name = generate_filename(title, '.csv')
+                self.report_file.save(
+                    file_name, ContentFile(csv_file.getvalue().encode())
+                )
+            elif file_type == 'xlsx':
+                xlsx_file = data_export.list_to_xlsx_file(
+                    single_chunk, title, header, widths
+                )
+                file_name = generate_filename(title, '.xlsx')
+                self.report_file.save(file_name, ContentFile(xlsx_file.read()))
+            self.report_file_creation = datetime.datetime.today()
+            self.save()
+            return
         else:
-            xlsx_file = data_export.list_to_xlsx_file(objects_list, title,
-                                                      header, widths)
-            title = generate_filename(title, '.xlsx')
-            self.report_file.save(title, ContentFile(xlsx_file.getvalue()))
-        self.report_file_creation = datetime.datetime.today()
-        self.save()
-        if email_to:
-            for email in email_to:
-                self.email_report(email=email)
-        elif getattr(settings, 'REPORT_BUILDER_EMAIL_NOTIFICATION', False):
-            if user.email:
-                self.email_report(user=user)
+            zip_buffer = BytesIO()
+            with zipfile.ZipFile(zip_buffer, 'w') as zip_file:
+                for index, chunk in enumerate(chunks):
+                    chunk_title = f'{title}_part{index + 1}.{file_type}'
+                    if file_type == 'csv':
+                        csv_file = data_export.list_to_csv_file(
+                            chunk, chunk_title, header, widths
+                        )
+                        zip_file.writestr(chunk_title, csv_file.getvalue().encode())
+                    elif file_type == 'xlsx':
+                        xlsx_file = data_export.list_to_xlsx_file(
+                            chunk, chunk_title, header, widths
+                        )
+                        zip_file.writestr(chunk_title, xlsx_file.read())
+            zip_filename = f'{title}.zip'
+            self.report_file.save(zip_filename, ContentFile(zip_buffer.getvalue()))
+            self.report_file_creation = datetime.datetime.today()
+            self.save()
+            if email_to:
+                for email in email_to:
+                    self.email_report(email=email)
+            elif getattr(settings, 'REPORT_BUILDER_EMAIL_NOTIFICATION', False):
+                if user and user.email:
+                    self.email_report(user=user)
+
+    @staticmethod
+    def chunk_data(data, chunk_size):
+        for i in range(0, len(data), chunk_size):
+            yield data[i : i + chunk_size]
+
 
     def run_report(self, file_type, user=None, queryset=None, asynchronous=False, scheduled=False,
                    email_to: str = None):
@@ -478,12 +516,16 @@ def run_report(self, file_type, user=None, queryset=None, asynchronous=False, sc
         for field in display_fields:
             header.append(field.name)
             widths.append(field.width)
+
+        chunk_size = 1000000
+        chunks = list(self.chunk_data(objects_list, chunk_size))
+
         if scheduled:
-            self.async_report_save(objects_list, title, header, widths, file_type, email_to=email_to)
+            self.async_report_save(chunks, title, header, widths, user, file_type)
         elif asynchronous:
             if user is None:
                 raise Exception('Cannot run async report without a user')
-            self.async_report_save(objects_list, title, header, widths, user, file_type)
+            self.async_report_save(chunks, title, header, widths, user, file_type)
         else:
             if file_type == 'csv':
                 return data_export.list_to_csv_response(

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="django-report-builder",
-    version="6.5.7",
+    version="6.6.0",
     author="David Burke",
     author_email="david@burkesoftware.com",
     maintainer="Mustafa Ülker",