From fb80a9c92e054b7b14e18d842ae963f0418ec40c Mon Sep 17 00:00:00 2001 From: Katie Smith Date: Mon, 14 Oct 2019 16:43:37 +0100 Subject: [PATCH] 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 | 44 +++++++++++-------- ...t_notification_dao_delete_notifications.py | 30 +++++++++++++ 2 files changed, 56 insertions(+), 18 deletions(-) diff --git a/app/dao/notifications_dao.py b/app/dao/notifications_dao.py index 7901647a0..8b8c20cb8 100644 --- a/app/dao/notifications_dao.py +++ b/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): - notifications = db.session.query( +def insert_update_notification_history(notification_type, date_to_delete_from, service_id, query_limit=10000): + 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( - NotificationHistory.__table__.c, - notifications - ) + notifications_count = notification_query.count() - stmt = stmt.on_conflict_do_update( - constraint="notification_history_pkey", - set_={"notification_status": stmt.excluded.status, - "reference": stmt.excluded.reference, - "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() + while offset < notifications_count: + stmt = insert(NotificationHistory).from_select( + NotificationHistory.__table__.c, + notification_query.limit(query_limit).offset(offset) + ) + + stmt = stmt.on_conflict_do_update( + constraint="notification_history_pkey", + set_={ + "notification_status": stmt.excluded.status, + "reference": stmt.excluded.reference, + "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( diff --git a/tests/app/dao/notification_dao/test_notification_dao_delete_notifications.py b/tests/app/dao/notification_dao/test_notification_dao_delete_notifications.py index 8856defc8..648ab303c 100644 --- a/tests/app/dao/notification_dao/test_notification_dao_delete_notifications.py +++ b/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)