feat: message group id for nightly, letters PDF and research mode

Author: joybytes

app/celery/letters_pdf_tasks.py

@@ -92,7 +92,10 @@ def get_pdf_for_templated_letter(self, notification_id):
         encoded_data = signing.encode(letter_data)
 
         notify_celery.send_task(
-            name=TaskNames.CREATE_PDF_FOR_TEMPLATED_LETTER, args=(encoded_data,), queue=QueueNames.SANITISE_LETTERS
+            name=TaskNames.CREATE_PDF_FOR_TEMPLATED_LETTER,
+            args=(encoded_data,),
+            queue=QueueNames.SANITISE_LETTERS,
+            MessageGroupId=self.message_group_id if self.message_group_id is not None else str(notification.service_id),
         )
     except Exception as e:
         try:
@@ -294,6 +297,7 @@ def sanitise_letter(self, filename):
                 "allow_international_letters": notification.service.has_permission(INTERNATIONAL_LETTERS),
             },
             queue=QueueNames.SANITISE_LETTERS,
+            MessageGroupId=self.message_group_id,
         )
     except Exception:
         try:
@@ -541,8 +545,8 @@ def replay_letters_in_error(filename=None):
                 sanitise_letter.apply_async([filename], queue=QueueNames.LETTERS)
 
 
-@notify_celery.task(name="resanitise-pdf")
-def resanitise_pdf(notification_id):
+@notify_celery.task(bind=True, name="resanitise-pdf")
+def resanitise_pdf(self, notification_id):
     """
     `notification_id` is the notification id for a PDF letter which was either uploaded or sent using the API.
 
@@ -570,11 +574,12 @@ def resanitise_pdf(notification_id):
             "allow_international_letters": notification.service.has_permission(INTERNATIONAL_LETTERS),
         },
         queue=QueueNames.SANITISE_LETTERS,
+        MessageGroupId=self.message_group_id if self.message_group_id is not None else str(notification.service_id),
     )
 
 
-@notify_celery.task(name="resanitise-letter-attachment")
-def resanitise_letter_attachment(service_id, attachment_id, original_filename):
+@notify_celery.task(bind=True, name="resanitise-letter-attachment")
+def resanitise_letter_attachment(self, service_id, attachment_id, original_filename):
     """
     `service_id` is the service id for a PDF letter attachment/template.
     `attachment_id` is the attachment id for a PDF letter attachment which was uploaded for a template.
@@ -593,4 +598,5 @@ def resanitise_letter_attachment(service_id, attachment_id, original_filename):
             "original_filename": original_filename,
         },
         queue=QueueNames.SANITISE_LETTERS,
+        MessageGroupId=self.message_group_id if self.message_group_id is not None else str(service_id),
     )

app/celery/nightly_tasks.py

@@ -82,8 +82,16 @@ def _remove_csv_files(job_types):
 @notify_celery.task(name="archive-unsubscribe-requests")
 def archive_unsubscribe_requests():
     for service_id in get_service_ids_with_unsubscribe_requests():
-        archive_batched_unsubscribe_requests.apply_async(queue=QueueNames.REPORTING, args=[service_id])
-        archive_old_unsubscribe_requests.apply_async(queue=QueueNames.REPORTING, args=[service_id])
+        archive_batched_unsubscribe_requests.apply_async(
+            queue=QueueNames.REPORTING,
+            args=[service_id],
+            MessageGroupId=str(service_id),
+        )
+        archive_old_unsubscribe_requests.apply_async(
+            queue=QueueNames.REPORTING,
+            args=[service_id],
+            MessageGroupId=str(service_id),
+        )
 
 
 @notify_celery.task(name="archive-batched-unsubscribe-requests")
@@ -163,6 +171,7 @@ def _delete_notifications_older_than_retention_by_type(
                 "datetime_to_delete_before": day_to_delete_backwards_from,
             },
             countdown=(i / len(flexible_data_retention)) * stagger_total_period.seconds,
+            MessageGroupId=str(f.service_id),
         )
 
     seven_days_ago = get_london_midnight_in_utc(convert_utc_to_bst(datetime.utcnow()).date() - timedelta(days=7))
@@ -185,6 +194,7 @@ def _delete_notifications_older_than_retention_by_type(
                 "datetime_to_delete_before": seven_days_ago,
             },
             countdown=(i / len(service_ids_to_purge)) * stagger_total_period.seconds,
+            MessageGroupId=str(service_id),
         )
 
     extra = {
@@ -204,8 +214,8 @@ def _delete_notifications_older_than_retention_by_type(
     )
 
 
-@notify_celery.task(name="delete-notifications-for-service-and-type")
-def delete_notifications_for_service_and_type(service_id, notification_type, datetime_to_delete_before):
+@notify_celery.task(bind=True, name="delete-notifications-for-service-and-type")
+def delete_notifications_for_service_and_type(self, service_id, notification_type, datetime_to_delete_before):
     start = datetime.utcnow()
     num_deleted = move_notifications_to_notification_history(
         notification_type,
@@ -237,12 +247,14 @@ def delete_notifications_for_service_and_type(service_id, notification_type, dat
         delete_notifications_for_service_and_type.apply_async(
             args=(service_id, notification_type, datetime_to_delete_before),
             queue=QueueNames.REPORTING,
+            MessageGroupId=self.message_group_id,
         )
     else:
         # now we've deleted all the real notifications, clean up the test notifications
         delete_test_notifications_for_service_and_type.apply_async(
             args=(service_id, notification_type, datetime_to_delete_before),
             queue=QueueNames.REPORTING,
+            MessageGroupId=self.message_group_id,
         )
 
 
@@ -254,6 +266,7 @@ def delete_test_notifications_for_service_and_type(service_id, notification_type
         delete_test_notifications_for_service_and_type.apply_async(
             args=(service_id, notification_type, datetime_to_delete_before),
             queue=QueueNames.REPORTING,
+            MessageGroupId=str(service_id),
         )
 
 

app/celery/research_mode_tasks.py

@@ -54,15 +54,19 @@ def send_sms_response(provider, reference, to):
     make_request(SMS_TYPE, provider, body, headers)
 
 
-def send_email_response(reference, to):
+def send_email_response(reference, to, service_id):
     if to == perm_fail_email:
         body = ses_hard_bounce_callback(reference)
     elif to == temp_fail_email:
         body = ses_soft_bounce_callback(reference)
     else:
         body = ses_notification_callback(reference)
 
-    process_ses_results.apply_async([body], queue=QueueNames.RESEARCH_MODE)
+    process_ses_results.apply_async(
+        [body],
+        queue=QueueNames.RESEARCH_MODE,
+        MessageGroupId=str(service_id),
+    )
 
 
 def send_letter_response(notification_id: uuid.UUID, billable_units: int, postage: str):

app/delivery/send_to_providers.py

@@ -163,7 +163,7 @@ def send_email_to_provider(notification):
         if notification.key_type == KEY_TYPE_TEST:
             notification.reference = str(create_uuid())
             update_notification_to_sending(notification, provider)
-            send_email_response(notification.reference, notification.to)
+            send_email_response(notification.reference, notification.to, notification.service_id)
         else:
             email_sender_name = service.custom_email_sender_name or service.name
             from_address = (

tests/app/celery/test_letters_pdf_tasks.py

@@ -55,6 +55,7 @@
     create_notification,
     create_service,
 )
+from tests.conftest import _with_message_group_id
 
 
 def test_should_have_decorated_tasks_functions():
@@ -75,7 +76,8 @@ def test_get_pdf_for_templated_letter_happy_path(mocker, sample_letter_notificat
     mock_generate_letter_pdf_filename = mocker.patch(
         "app.celery.letters_pdf_tasks.generate_letter_pdf_filename", return_value="LETTER.PDF"
     )
-    get_pdf_for_templated_letter(sample_letter_notification.id)
+    with _with_message_group_id(get_pdf_for_templated_letter, str(sample_letter_notification.service_id)):
+        get_pdf_for_templated_letter(sample_letter_notification.id)
 
     letter_data = {
         "letter_contact_block": sample_letter_notification.reply_to_text,
@@ -97,7 +99,10 @@ def test_get_pdf_for_templated_letter_happy_path(mocker, sample_letter_notificat
     }
 
     mock_celery.assert_called_once_with(
-        name=TaskNames.CREATE_PDF_FOR_TEMPLATED_LETTER, args=(ANY,), queue=QueueNames.SANITISE_LETTERS
+        name=TaskNames.CREATE_PDF_FOR_TEMPLATED_LETTER,
+        args=(ANY,),
+        queue=QueueNames.SANITISE_LETTERS,
+        MessageGroupId=str(sample_letter_notification.service_id),
     )
 
     actual_data = signing.decode(mock_celery.call_args.kwargs["args"][0])
@@ -117,7 +122,8 @@ def test_get_pdf_for_templated_letter_with_letter_attachment(mocker, sample_lett
 
     mock_celery = mocker.patch("app.celery.letters_pdf_tasks.notify_celery.send_task")
     mocker.patch("app.celery.letters_pdf_tasks.generate_letter_pdf_filename", return_value="LETTER.PDF")
-    get_pdf_for_templated_letter(sample_letter_notification.id)
+    with _with_message_group_id(get_pdf_for_templated_letter, str(sample_letter_notification.service_id)):
+        get_pdf_for_templated_letter(sample_letter_notification.id)
 
     actual_data = signing.decode(mock_celery.call_args.kwargs["args"][0])
 
@@ -134,8 +140,9 @@ def test_get_pdf_for_templated_letter_retries_upon_error(mocker, sample_letter_n
     mocker.patch("app.celery.letters_pdf_tasks.generate_letter_pdf_filename", return_value="LETTER.PDF")
     mock_retry = mocker.patch("app.celery.letters_pdf_tasks.get_pdf_for_templated_letter.retry")
 
-    with caplog.at_level("ERROR"):
-        get_pdf_for_templated_letter(sample_letter_notification.id)
+    with _with_message_group_id(get_pdf_for_templated_letter, str(sample_letter_notification.service_id)):
+        with caplog.at_level("ERROR"):
+            get_pdf_for_templated_letter(sample_letter_notification.id)
 
     assert mock_celery.called
     assert mock_retry.called
@@ -153,8 +160,9 @@ def test_get_pdf_for_templated_letter_sets_technical_failure_max_retries(mocker,
     )
     mock_update_noti = mocker.patch("app.celery.letters_pdf_tasks.update_notification_status_by_id")
 
-    with pytest.raises(NotificationTechnicalFailureException) as e:
-        get_pdf_for_templated_letter(sample_letter_notification.id)
+    with _with_message_group_id(get_pdf_for_templated_letter, str(sample_letter_notification.service_id)):
+        with pytest.raises(NotificationTechnicalFailureException) as e:
+            get_pdf_for_templated_letter(sample_letter_notification.id)
 
     assert (
         e.value.args[0] == f"RETRY FAILED: Max retries reached. "
@@ -390,9 +398,17 @@ def test_send_letters_volume_email_to_dvla(notify_db_session, mock_celery_task,
 
     emails_to_dvla = Notification.query.all()
     assert len(emails_to_dvla) == 2
-    send_mock.called = 2
-    send_mock.assert_any_call([str(emails_to_dvla[0].id)], queue=QueueNames.NOTIFY)
-    send_mock.assert_any_call([str(emails_to_dvla[1].id)], queue=QueueNames.NOTIFY)
+    assert send_mock.call_count == 2
+    send_mock.assert_any_call(
+        [str(emails_to_dvla[0].id)],
+        queue=QueueNames.NOTIFY,
+        MessageGroupId=str(emails_to_dvla[0].service_id),
+    )
+    send_mock.assert_any_call(
+        [str(emails_to_dvla[1].id)],
+        queue=QueueNames.NOTIFY,
+        MessageGroupId=str(emails_to_dvla[1].service_id),
+    )
     for email in emails_to_dvla:
         assert str(email.template_id) == current_app.config["LETTERS_VOLUME_EMAIL_TEMPLATE_ID"]
         assert email.to in current_app.config["DVLA_EMAIL_ADDRESSES"]
@@ -451,7 +467,8 @@ def test_sanitise_letter_calls_template_preview_sanitise_task(
     sample_letter_notification.service = create_service(service_permissions=permissions)
     sample_letter_notification.status = NOTIFICATION_PENDING_VIRUS_CHECK
 
-    sanitise_letter(filename)
+    with _with_message_group_id(sanitise_letter, str(sample_letter_notification.service_id)):
+        sanitise_letter(filename)
 
     mock_celery.assert_called_once_with(
         name=TaskNames.SANITISE_LETTER,
@@ -461,6 +478,7 @@ def test_sanitise_letter_calls_template_preview_sanitise_task(
             "allow_international_letters": expected_international_letters_allowed,
         },
         queue=QueueNames.SANITISE_LETTERS,
+        MessageGroupId=str(sample_letter_notification.service_id),
     )
 
 
@@ -881,7 +899,8 @@ def test_resanitise_pdf_calls_template_preview_with_letter_details(
     sample_letter_notification.created_at = datetime(2021, 2, 7, 12)
     sample_letter_notification.service = create_service(service_permissions=permissions)
 
-    resanitise_pdf(sample_letter_notification.id)
+    with _with_message_group_id(resanitise_pdf, str(sample_letter_notification.service_id)):
+        resanitise_pdf(sample_letter_notification.id)
 
     mock_celery.assert_called_once_with(
         name=TaskNames.RECREATE_PDF_FOR_PRECOMPILED_LETTER,
@@ -904,7 +923,8 @@ def test_resanitise_letter_attachment_calls_template_preview_with_attachment_det
     attachment_id = str(uuid.uuid4())
     original_filename = "test-123abc.pdf"
 
-    resanitise_letter_attachment(service_id, attachment_id, original_filename)
+    with _with_message_group_id(resanitise_letter_attachment, service_id):
+        resanitise_letter_attachment(service_id, attachment_id, original_filename)
 
     mock_celery.assert_called_once_with(
         name=TaskNames.RECREATE_PDF_FOR_TEMPLATE_LETTER_ATTACHMENTS,
@@ -914,4 +934,5 @@ def test_resanitise_letter_attachment_calls_template_preview_with_attachment_det
             "original_filename": original_filename,
         },
         queue=QueueNames.SANITISE_LETTERS,
+        MessageGroupId=service_id,
     )

tests/app/celery/test_nightly_tasks.py

@@ -63,6 +63,7 @@
     create_unsubscribe_request,
     create_unsubscribe_request_report,
 )
+from tests.conftest import _with_message_group_id
 
 
 @freeze_time("2016-10-18T10:00:00")
@@ -175,6 +176,12 @@ def test_archive_unsubscribe_requests(notify_db_session, mock_celery_task):
         == {service.id for service in services_with_requests}
     )
 
+    assert (
+        {call[1]["MessageGroupId"] for call in mock_archive_processed.call_args_list}
+        == {call[1]["MessageGroupId"] for call in mock_archive_old.call_args_list}
+        == {str(service.id) for service in services_with_requests}
+    )
+
     assert (
         [call[1]["queue"] for call in mock_archive_processed.call_args_list]
         == [call[1]["queue"] for call in mock_archive_old.call_args_list]
@@ -570,6 +577,7 @@ def test_delete_notifications_task_calls_task_for_services_with_data_retention_o
             "datetime_to_delete_before": datetime(2021, 6, 1, 23, 0),
         },
         countdown=0.0,
+        MessageGroupId=str(sms_service.id),
     )
 
 
@@ -598,6 +606,7 @@ def test_delete_notifications_task_calls_task_for_services_with_data_retention_b
                     "datetime_to_delete_before": datetime(2021, 3, 22, 0, 0),
                 },
                 countdown=ANY,
+                MessageGroupId=str(service_14_days.id),
             ),
             call(
                 queue=ANY,
@@ -607,6 +616,7 @@ def test_delete_notifications_task_calls_task_for_services_with_data_retention_b
                     "datetime_to_delete_before": datetime(2021, 4, 1, 23, 0),
                 },
                 countdown=ANY,
+                MessageGroupId=str(service_3_days.id),
             ),
         ],
     )
@@ -655,6 +665,7 @@ def test_delete_notifications_task_calls_task_for_services_that_have_sent_notifi
                     "datetime_to_delete_before": datetime(2021, 3, 27, 0, 0),
                 },
                 countdown=ANY,
+                MessageGroupId=str(service_will_delete_1.id),
             ),
             call(
                 queue=ANY,
@@ -664,6 +675,7 @@ def test_delete_notifications_task_calls_task_for_services_that_have_sent_notifi
                     "datetime_to_delete_before": datetime(2021, 3, 27, 0, 0),
                 },
                 countdown=ANY,
+                MessageGroupId=str(service_will_delete_2.id),
             ),
         ],
     )
@@ -682,12 +694,15 @@ def test_delete_notifications_for_service_and_type_queues_up_second_task_if_thin
     notification_type = "some-str"
     datetime_to_delete_before = datetime.utcnow()
 
-    delete_notifications_for_service_and_type(service_id, notification_type, datetime_to_delete_before)
+    with _with_message_group_id(delete_notifications_for_service_and_type, str(service_id)):
+        delete_notifications_for_service_and_type(service_id, notification_type, datetime_to_delete_before)
 
     mock_move.assert_called_once_with(notification_type, service_id, datetime_to_delete_before)
     # the next task is queued up with the exact same args
     mock_task_call.assert_called_once_with(
-        args=(service_id, notification_type, datetime_to_delete_before), queue="reporting-tasks"
+        args=(service_id, notification_type, datetime_to_delete_before),
+        queue="reporting-tasks",
+        MessageGroupId=str(service_id),
     )
     assert not mock_delete_tests.called
 
@@ -703,13 +718,16 @@ def test_delete_notifications_for_service_and_type_removes_test_notifications_if
     notification_type = "some-str"
     datetime_to_delete_before = datetime.utcnow()
 
-    delete_notifications_for_service_and_type(service_id, notification_type, datetime_to_delete_before)
+    with _with_message_group_id(delete_notifications_for_service_and_type, str(service_id)):
+        delete_notifications_for_service_and_type(service_id, notification_type, datetime_to_delete_before)
 
     mock_move.assert_called_once_with(notification_type, service_id, datetime_to_delete_before)
     # the next task is not queued up
     assert not mock_delete_live_notis_task_call.called
     mock_delete_tests_task_call.assert_called_once_with(
-        args=(service_id, notification_type, datetime_to_delete_before), queue="reporting-tasks"
+        args=(service_id, notification_type, datetime_to_delete_before),
+        queue="reporting-tasks",
+        MessageGroupId=str(service_id),
     )
 
 
@@ -727,7 +745,9 @@ def test_delete_test_notifications_for_service_and_type_queues_up_second_task_if
 
     mock_delete.assert_called_once_with(notification_type, service_id, datetime_to_delete_before)
     mock_task_call.assert_called_once_with(
-        args=(service_id, notification_type, datetime_to_delete_before), queue="reporting-tasks"
+        args=(service_id, notification_type, datetime_to_delete_before),
+        queue="reporting-tasks",
+        MessageGroupId=str(service_id),
     )