Refactor get notifications streaming including formatting

app/__init__.py

@@ -311,7 +311,7 @@ def format_notification_status(status, template_type):
             'sending': 'Sending',
             'created': 'Sending'
         }
-    }.get(template_type).get(status, status)
+    }[template_type].get(status, status)
 
 
 def format_notification_status_as_time(status, created, updated):

app/main/views/jobs.py

@@ -104,7 +104,6 @@ def view_job(service_id, job_id):
 
     total_notifications = job.get('notification_count', 0)
     processed_notifications = job.get('notifications_delivered', 0) + job.get('notifications_failed', 0)
-
     return render_template(
         'views/jobs/job.html',
         finished=(total_notifications == processed_notifications),
@@ -148,7 +147,8 @@ def view_job_csv(service_id, job_id):
             service_id=service_id,
             job_id=job_id,
             status=filter_args.get('status'),
-            page_size=job['notification_count']
+            page=request.args.get('page', 1),
+            page_size=5000
         ),
         mimetype='text/csv',
         headers={
@@ -205,9 +205,22 @@ def get_notifications(service_id, message_type, status_override=None):
         abort(404, "Invalid page argument ({}) reverting to page 1.".format(request.args['page'], None))
     if message_type not in ['email', 'sms']:
         abort(404)
-
     filter_args = _parse_filter_args(request.args)
     filter_args['status'] = _set_status_filters(filter_args)
+    if request.path.endswith('csv'):
+        return Response(
+            generate_notifications_csv(
+                service_id=service_id,
+                page=page,
+                page_size=5000,
+                template_type=[message_type],
+                status=filter_args.get('status'),
+                limit_days=current_app.config['ACTIVITY_STATS_LIMIT_DAYS']
+            ),
+            mimetype='text/csv',
+            headers={
+                'Content-Disposition': 'inline; filename="notifications.csv"'}
+        )
 
     notifications = notification_api_client.get_notifications_for_service(
         service_id=service_id,
@@ -221,27 +234,14 @@ def get_notifications(service_id, message_type, status_override=None):
         'status': request.args.get('status')
     }
     prev_page = None
+
     if notifications['links'].get('prev', None):
         prev_page = generate_previous_dict('main.view_notifications', service_id, page, url_args=url_args)
     next_page = None
+
     if notifications['links'].get('next', None):
         next_page = generate_next_dict('main.view_notifications', service_id, page, url_args)
 
-    if request.path.endswith('csv'):
-        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']
-            ),
-            mimetype='text/csv',
-            headers={
-                'Content-Disposition': 'inline; filename="notifications.csv"'}
-        )
-
     return {
         'counts': render_template(
             'views/activity/counts.html',

app/utils.py

@@ -115,9 +115,14 @@ def get_errors_for_csv(recipients, template_type):
     return errors
 
 
-def generate_notifications_csv(*args, **kwargs):
+def generate_notifications_csv(**kwargs):
+    from app import notification_api_client
+
     csvfile = StringIO()
-    csvwriter = csv.writer(csvfile)
+    csvwriter = csv.DictWriter(
+        csvfile,
+        fieldnames=['Row number', 'Recipient', 'Template', 'Type', 'Job', 'Status', 'Time']
+    )
 
     def read_current_row():
         csvfile.seek(0)
@@ -129,27 +134,42 @@ def generate_notifications_csv(*args, **kwargs):
         csvfile.truncate()
 
     # Initiate download quicker by returning the headers first
-    csvwriter.writerow(['Row number', 'Recipient', 'Template', 'Type', 'Job', 'Status', 'Time'])
+    csvwriter.writeheader()
     line = read_current_row()
     clear_file_buffer()
     yield line
 
-    from app import format_datetime_24h, format_notification_status, notification_api_client
+    # get_notifications_for_service is always paginated by the api, so set a page param if not specified
-    json_list = notification_api_client.get_notifications_for_service(**kwargs)['notifications']
+    if 'page' not in kwargs:
-    for x in json_list:
+        kwargs['page'] = 1
-        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()
+    while kwargs['page']:
-        clear_file_buffer()
+        notifications_resp = notification_api_client.get_notifications_for_service(**kwargs)
-        yield line
+        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
+    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():