From 8cf8d24e3717973fd29456a2578efd9fcad7f52e Mon Sep 17 00:00:00 2001
From: Pea Tyczynska <pea.tyczynska@digital.cabinet-office.gov.uk>
Date: Tue, 24 Sep 2019 16:52:18 +0100
Subject: [PATCH] Return count of notifications in the database for a job

When we cancel a job, we need to check if all notifications are
already in the database. So far, we were querying for all
notification objects in the database and counting them in
admin app, which runs into pagination problems for large jobs,
and could time out for very large jobs.
---
 app/dao/notifications_dao.py                  |  5 +++++
 app/job/rest.py                               | 10 ++++++++-
 .../notification_dao/test_notification_dao.py | 22 +++++++++++++++++++
 tests/app/job/test_rest.py                    |  8 +++++++
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/app/dao/notifications_dao.py b/app/dao/notifications_dao.py
index 8976bbebc..7c4e1e91a 100644
--- a/app/dao/notifications_dao.py
+++ b/app/dao/notifications_dao.py
@@ -178,6 +178,11 @@ def get_notifications_for_job(service_id, job_id, filter_dict=None, page=1, page
     )
 
 
+@statsd(namespace="dao")
+def dao_get_notification_count_for_job_id(job_id):
+    return Notification.query.filter_by(job_id=job_id).count()
+
+
 @statsd(namespace="dao")
 def get_notification_with_personalisation(service_id, notification_id, key_type):
     filter_dict = {'service_id': service_id, 'id': notification_id}
diff --git a/app/job/rest.py b/app/job/rest.py
index 145f3c004..e2abe2bcb 100644
--- a/app/job/rest.py
+++ b/app/job/rest.py
@@ -20,7 +20,7 @@ from app.dao.jobs_dao import (
 from app.dao.fact_notification_status_dao import fetch_notification_statuses_for_job
 from app.dao.services_dao import dao_fetch_service_by_id
 from app.dao.templates_dao import dao_get_template_by_id
-from app.dao.notifications_dao import get_notifications_for_job
+from app.dao.notifications_dao import dao_get_notification_count_for_job_id, get_notifications_for_job
 from app.schemas import (
     job_schema,
     unarchived_template_schema,
@@ -108,6 +108,14 @@ def get_all_notifications_for_service_job(service_id, job_id):
     ), 200
 
 
+@job_blueprint.route('/<job_id>/notification_count', methods=['GET'])
+def get_notification_count_for_job_id(service_id, job_id):
+    count = dao_get_notification_count_for_job_id(job_id)
+    return jsonify(
+        count=count
+    ), 200
+
+
 @job_blueprint.route('', methods=['GET'])
 def get_jobs_by_service(service_id):
     if request.args.get('limit_days'):
diff --git a/tests/app/dao/notification_dao/test_notification_dao.py b/tests/app/dao/notification_dao/test_notification_dao.py
index 8bbbb98bd..2cfa369d7 100644
--- a/tests/app/dao/notification_dao/test_notification_dao.py
+++ b/tests/app/dao/notification_dao/test_notification_dao.py
@@ -15,6 +15,7 @@ from app.dao.notifications_dao import (
     dao_get_last_notification_added_for_job_id,
     dao_get_last_template_usage,
     dao_get_notifications_by_to_field,
+    dao_get_notification_count_for_job_id,
     dao_get_scheduled_notifications,
     dao_timeout_notifications,
     dao_update_notification,
@@ -554,6 +555,27 @@ def test_get_all_notifications_for_job_by_status(sample_job):
     assert len(notifications(filter_dict={'status': NOTIFICATION_STATUS_TYPES[:3]}).items) == 3
 
 
+def test_dao_get_notification_count_for_job_id(notify_db_session, notify_db):
+    service = create_service()
+    template = create_template(service)
+    job = create_job(template, notification_count=3)
+    for i in range(3):
+        create_notification(job=job)
+
+    create_notification(template)
+
+    assert dao_get_notification_count_for_job_id(job.id) == 3
+
+
+def test_dao_get_notification_count_for_job_id_only_finds_notification_already_in_db(notify_db_session, notify_db):
+    service = create_service()
+    template = create_template(service)
+    job = create_job(template, notification_count=3)
+    create_notification(template)
+
+    assert dao_get_notification_count_for_job_id(job.id) == 0
+
+
 def test_update_notification_sets_status(sample_notification):
     assert sample_notification.status == 'created'
     sample_notification.status = 'failed'
diff --git a/tests/app/job/test_rest.py b/tests/app/job/test_rest.py
index d81210491..312e46e92 100644
--- a/tests/app/job/test_rest.py
+++ b/tests/app/job/test_rest.py
@@ -563,6 +563,14 @@ def test_get_all_notifications_for_job_returns_correct_format(
     assert resp['notifications'][0]['status'] == sample_notification_with_job.status
 
 
+def test_get_notification_count_for_job_id(admin_request, mocker):
+    mock_dao = mocker.patch('app.job.rest.dao_get_notification_count_for_job_id', return_value=3)
+    job_id = uuid.uuid4()
+    response = admin_request.get('job.get_notification_count_for_job_id', service_id=uuid.uuid4(), job_id=job_id)
+    mock_dao.assert_called_once_with(str(job_id))
+    assert response["count"] == 3
+
+
 def test_get_job_by_id(admin_request, sample_job):
     job_id = str(sample_job.id)
     service_id = sample_job.service.id