Merge pull request #1070 from alphagov/return-streamed-csv

Fix timeout on download of large CSV reports

app/utils.py

@@ -5,7 +5,14 @@ from os import path
 from functools import wraps
 import unicodedata
 
-from flask import (abort, current_app, session, request, redirect, url_for)
+from flask import (
+    abort,
+    current_app,
+    redirect,
+    request,
+    session,
+    url_for
+)
 from flask_login import current_user
 
 from wand.image import Image
@@ -111,24 +118,61 @@ def get_errors_for_csv(recipients, template_type):
     return errors
 
 
-def generate_notifications_csv(json_list):
+def generate_notifications_csv(**kwargs):
+    from app import notification_api_client
+
+    csvfile = StringIO()
+    csvwriter = csv.DictWriter(
+        csvfile,
+        fieldnames=['Row number', 'Recipient', 'Template', 'Type', 'Job', 'Status', 'Time']
+    )
+
+    def read_current_row():
+        csvfile.seek(0)
+        line = csvfile.read()
+        return line
+
+    def clear_file_buffer():
+        csvfile.seek(0)
+        csvfile.truncate()
+
+    # Initiate download quicker by returning the headers first
+    csvwriter.writeheader()
+    line = read_current_row()
+    clear_file_buffer()
+    yield line
+
+    # get_notifications_for_service is always paginated by the api, so set a page param if not specified
+    if 'page' not in kwargs:
+        kwargs['page'] = 1
+
+    while kwargs['page']:
+        notifications_resp = notification_api_client.get_notifications_for_service(**kwargs)
+        notifications_list = notifications_resp['notifications']
+        for x in notifications_list:
+            csvwriter.writerow(format_notification_for_csv(x))
+
+            line = read_current_row()
+            clear_file_buffer()
+            yield line
+
+        if notifications_resp['links'].get('next'):
+            kwargs['page'] += 1
+        else:
+            return
+
+
+def format_notification_for_csv(notification):
     from app import format_datetime_24h, format_notification_status
-    content = StringIO()
-    retval = None
-    with content as csvfile:
-        csvwriter = csv.writer(csvfile)
-        csvwriter.writerow(['Row number', 'Recipient', 'Template', 'Type', 'Job', 'Status', 'Time'])
-        for x in json_list:
-            csvwriter.writerow([
-                int(x['job_row_number']) + 2 if 'job_row_number' in x and x['job_row_number'] else '',
-                x['to'],
-                x['template']['name'],
-                x['template']['template_type'],
-                x['job']['original_file_name'] if x['job'] else '',
-                format_notification_status(x['status'], x['template']['template_type']),
-                format_datetime_24h(x['created_at'])])
-        retval = content.getvalue()
-    return retval
+    return {
+        'Row number': int(notification['job_row_number']) + 2 if notification.get('job_row_number') else '',
+        'Recipient': notification['to'],
+        'Template': notification['template']['name'],
+        'Type': notification['template']['template_type'],
+        'Job': notification['job']['original_file_name'] if notification['job'] else '',
+        'Status': format_notification_status(notification['status'], notification['template']['template_type']),
+        'Time': format_datetime_24h(notification['created_at'])
+    }
 
 
 def get_page_from_request():