Return csv download stream instead of all-in-memory

app/main/views/jobs.py

@@ -12,7 +12,8 @@ from flask import (
     request,
     url_for,
     current_app,
-    redirect
+    redirect,
+    Response
 )
 from flask_login import login_required
 from werkzeug.datastructures import MultiDict
@@ -142,23 +143,19 @@ def view_job_csv(service_id, job_id):
     filter_args = _parse_filter_args(request.args)
     filter_args['status'] = _set_status_filters(filter_args)
 
-    return (
+    return Response(
         generate_notifications_csv(
-            notification_api_client.get_notifications_for_service(
-                service_id,
-                job_id,
-                status=filter_args.get('status'),
-                page_size=job['notification_count']
-            )['notifications']
+            service_id=service_id,
+            job_id=job_id,
+            status=filter_args.get('status'),
+            page_size=job['notification_count']
         ),
-        200,
-        {
-            'Content-Type': 'text/csv; charset=utf-8',
+        mimetype='text/csv',
+        headers={
             'Content-Disposition': 'inline; filename="{} - {}.csv"'.format(
                 template['name'],
                 format_datetime_short(job['created_at'])
-            )
-        }
+            )}
     )
 
 
@@ -231,18 +228,20 @@ def get_notifications(service_id, message_type, status_override=None):
         next_page = generate_next_dict('main.view_notifications', service_id, page, url_args)
 
     if request.path.endswith('csv'):
-        csv_content = generate_notifications_csv(
-            notification_api_client.get_notifications_for_service(
+        return Response(
+            generate_notifications_csv(
                 service_id=service_id,
                 page=page,
                 page_size=notifications['total'],
                 template_type=[message_type],
                 status=filter_args.get('status'),
-                limit_days=current_app.config['ACTIVITY_STATS_LIMIT_DAYS'])['notifications'])
-        return csv_content, 200, {
-            'Content-Type': 'text/csv; charset=utf-8',
-            'Content-Disposition': 'inline; filename="notifications.csv"'
-        }
+                limit_days=current_app.config['ACTIVITY_STATS_LIMIT_DAYS']
+            ),
+            mimetype='text/csv',
+            headers={
+                'Content-Disposition': 'inline; filename="notifications.csv"'}
+        )
+
     return {
         'counts': render_template(
             'views/activity/counts.html',
@@ -296,7 +295,7 @@ def get_status_filters(service, message_type, statistics):
             stats[key]
         )
         for key, label, option in filters
-        ]
+    ]
 
 
 def _get_job_counts(job, help_argument):

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 notifications_utils.template import (
@@ -108,24 +115,41 @@ def get_errors_for_csv(recipients, template_type):
     return errors
 
 
-def generate_notifications_csv(json_list):
-    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
+def generate_notifications_csv(*args, **kwargs):
+    csvfile = StringIO()
+    csvwriter = csv.writer(csvfile)
+
+    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.writerow(['Row number', 'Recipient', 'Template', 'Type', 'Job', 'Status', 'Time'])
+    line = read_current_row()
+    clear_file_buffer()
+    yield line
+
+    from app import format_datetime_24h, format_notification_status, notification_api_client
+    json_list = notification_api_client.get_notifications_for_service(**kwargs)['notifications']
+    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'])
+        ])
+
+        line = read_current_row()
+        clear_file_buffer()
+        yield line
 
 
 def get_page_from_request():