Update insert_update_notification_history to take a query limit

The nightly job to delete email notifications was failing because it was
timing out (`psycopg2.errors.QueryCanceled: canceling statement due to statement timeout`).

This adds a query limit to the query which inserts or updates
notification history so that it only updates a maximum of 10000 rows at
a time.

app/dao/notifications_dao.py

@@ -367,8 +367,10 @@ def _delete_for_query(subquery):
     return deleted
 
 
-def insert_update_notification_history(notification_type, date_to_delete_from, service_id):
+def insert_update_notification_history(notification_type, date_to_delete_from, service_id, query_limit=10000):
-    notifications = db.session.query(
+    offset = 0
+
+    notification_query = db.session.query(
         *[x.name for x in NotificationHistory.__table__.c]
     ).filter(
         Notification.notification_type == notification_type,
@@ -376,23 +378,29 @@ def insert_update_notification_history(notification_type, date_to_delete_from, s
         Notification.created_at < date_to_delete_from,
         Notification.key_type != KEY_TYPE_TEST
     )
-    stmt = insert(NotificationHistory).from_select(
+    notifications_count = notification_query.count()
-        NotificationHistory.__table__.c,
-        notifications
-    )
 
-    stmt = stmt.on_conflict_do_update(
+    while offset < notifications_count:
-        constraint="notification_history_pkey",
+        stmt = insert(NotificationHistory).from_select(
-        set_={"notification_status": stmt.excluded.status,
+            NotificationHistory.__table__.c,
-              "reference": stmt.excluded.reference,
+            notification_query.limit(query_limit).offset(offset)
-              "billable_units": stmt.excluded.billable_units,
+        )
-              "updated_at": stmt.excluded.updated_at,
+
-              "sent_at": stmt.excluded.sent_at,
+        stmt = stmt.on_conflict_do_update(
-              "sent_by": stmt.excluded.sent_by
+            constraint="notification_history_pkey",
-              }
+            set_={
-    )
+                "notification_status": stmt.excluded.status,
-    db.session.connection().execute(stmt)
+                "reference": stmt.excluded.reference,
-    db.session.commit()
+                "billable_units": stmt.excluded.billable_units,
+                "updated_at": stmt.excluded.updated_at,
+                "sent_at": stmt.excluded.sent_at,
+                "sent_by": stmt.excluded.sent_by
+            }
+        )
+        db.session.connection().execute(stmt)
+        db.session.commit()
+
+        offset += query_limit
 
 
 def _delete_letters_from_s3(

tests/app/dao/notification_dao/test_notification_dao_delete_notifications.py

@@ -9,6 +9,7 @@ from freezegun import freeze_time
 
 from app.dao.notifications_dao import (
     delete_notifications_older_than_retention_by_type,
+    db,
     insert_update_notification_history
 )
 from app.models import Notification, NotificationHistory
@@ -287,6 +288,35 @@ def test_insert_update_notification_history(sample_service):
     assert notification_3.id in history_ids
 
 
+def test_insert_update_notification_history_with_more_notifications_than_query_limit(mocker, sample_service):
+    template = create_template(sample_service, template_type='sms')
+    notification_1 = create_notification(template=template, created_at=datetime.utcnow() - timedelta(days=3))
+    notification_2 = create_notification(template=template, created_at=datetime.utcnow() - timedelta(days=8))
+    notification_3 = create_notification(template=template, created_at=datetime.utcnow() - timedelta(days=9))
+    other_types = ['email', 'letter']
+    for template_type in other_types:
+        t = create_template(service=sample_service, template_type=template_type)
+        create_notification(template=t, created_at=datetime.utcnow() - timedelta(days=3))
+        create_notification(template=t, created_at=datetime.utcnow() - timedelta(days=8))
+
+    db_connection_spy = mocker.spy(db.session, 'connection')
+    db_commit_spy = mocker.spy(db.session, 'commit')
+
+    insert_update_notification_history(
+        notification_type='sms', date_to_delete_from=datetime.utcnow() - timedelta(days=7),
+        service_id=sample_service.id, query_limit=1)
+    history = NotificationHistory.query.all()
+    assert len(history) == 2
+
+    history_ids = [x.id for x in history]
+    assert notification_1.id not in history_ids
+    assert notification_2.id in history_ids
+    assert notification_3.id in history_ids
+
+    assert db_connection_spy.call_count == 2
+    assert db_commit_spy.call_count == 2
+
+
 def test_insert_update_notification_history_only_insert_update_given_service(sample_service):
     other_service = create_service(service_name='another service')
     other_template = create_template(service=other_service)