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

@@ -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}

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'):

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'

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