feat: message group id for job pipeline (bulk upload, tasks, scheduled)

Author: joybytes

app/celery/scheduled_tasks.py

@@ -96,7 +96,11 @@
 def run_scheduled_jobs():
     try:
         for job in dao_set_scheduled_jobs_to_pending():
-            process_job.apply_async([str(job.id)], queue=QueueNames.JOBS)
+            process_job.apply_async(
+                [str(job.id)],
+                queue=QueueNames.JOBS,
+                MessageGroupId=str(job.service_id),
+            )
             current_app.logger.info("Job ID %s added to process job queue", job.id, extra={"job_id": job.id})
     except SQLAlchemyError:
         current_app.logger.exception("Failed to run scheduled jobs")
@@ -335,7 +339,11 @@ def replay_created_notifications() -> None:
             letter.id,
             extra={"notification_id": letter.id},
         )
-        get_pdf_for_templated_letter.apply_async([str(letter.id)], queue=QueueNames.CREATE_LETTERS_PDF)
+        get_pdf_for_templated_letter.apply_async(
+            [str(letter.id)],
+            queue=QueueNames.CREATE_LETTERS_PDF,
+            MessageGroupId=str(letter.service_id),
+        )
 
 
 @notify_celery.task(name="check-if-letters-still-pending-virus-check")
@@ -359,6 +367,7 @@ def check_if_letters_still_pending_virus_check(max_minutes_ago_to_check: int = 3
                 name=TaskNames.SCAN_FILE,
                 kwargs={"filename": filename},
                 queue=QueueNames.ANTIVIRUS,
+                MessageGroupId=str(letter.service_id),
             )
         else:
             current_app.logger.warning(
@@ -438,7 +447,7 @@ def check_for_missing_rows_in_completed_jobs():
 
             extra = {"job_row_number": row_to_process.missing_row, "job_id": job.id}
             current_app.logger.info("Processing missing row %(job_row_number)s for job %(job_id)s", extra, extra=extra)
-            process_job_row(template.template_type, task_args_kwargs)
+            process_job_row(template.template_type, task_args_kwargs, str(job.service_id))
 
 
 @notify_celery.task(name="update-status-of-fully-processed-jobs")

app/celery/tasks.py

@@ -68,8 +68,8 @@ class ProcessReportRequestException(Exception):
     pass
 
 
-@notify_celery.task(name="process-job")
-def process_job(job_id, sender_id=None, shatter_batch_size=DEFAULT_SHATTER_JOB_ROWS_BATCH_SIZE):
+@notify_celery.task(bind=True, name="process-job")
+def process_job(self, job_id, sender_id=None, shatter_batch_size=DEFAULT_SHATTER_JOB_ROWS_BATCH_SIZE):
     start = datetime.utcnow()
     job = dao_get_job_by_id(job_id)
     current_app.logger.info(
@@ -116,19 +116,20 @@ def process_job(job_id, sender_id=None, shatter_batch_size=DEFAULT_SHATTER_JOB_R
             get_id_task_args_kwargs_for_job_row(row, template, job, service, sender_id=sender_id)[1]
             for row in shatter_batch
         ]
-        _shatter_job_rows_with_subdivision(template.template_type, batch_args_kwargs)
+        _shatter_job_rows_with_subdivision(template.template_type, batch_args_kwargs, self.message_group_id)
 
     job_complete(job, start=start)
 
 
-def _shatter_job_rows_with_subdivision(template_type, args_kwargs_seq, top_level=True):
+def _shatter_job_rows_with_subdivision(template_type, args_kwargs_seq, message_group_id, top_level=True):
     try:
         shatter_job_rows.apply_async(
             (
                 template_type,
                 args_kwargs_seq,
             ),
             queue=QueueNames.JOBS,
+            MessageGroupId=message_group_id,
         )
     except BotoClientError as e:
         # this information is helpfully not preserved outside the message string of the exception, so
@@ -146,7 +147,7 @@ def _shatter_job_rows_with_subdivision(template_type, args_kwargs_seq, top_level
             raise UnprocessableJobRow from e
 
         for sub_batch in (args_kwargs_seq[:split_batch_size], args_kwargs_seq[split_batch_size:]):
-            _shatter_job_rows_with_subdivision(template_type, sub_batch, top_level=False)
+            _shatter_job_rows_with_subdivision(template_type, sub_batch, message_group_id, top_level=False)
 
     else:
         if not top_level:
@@ -156,16 +157,17 @@ def _shatter_job_rows_with_subdivision(template_type, args_kwargs_seq, top_level
             )
 
 
-@notify_celery.task(name="shatter-job-rows")
+@notify_celery.task(bind=True, name="shatter-job-rows")
 def shatter_job_rows(
+    self,
     template_type: str,
     args_kwargs_seq: Sequence,
 ):
     for task_args_kwargs in args_kwargs_seq:
-        process_job_row(template_type, task_args_kwargs)
+        process_job_row(template_type, task_args_kwargs, self.message_group_id)
 
 
-def process_job_row(template_type, task_args_kwargs):
+def process_job_row(template_type, task_args_kwargs, message_group_id=None):
     send_fn = {
         SMS_TYPE: save_sms,
         EMAIL_TYPE: save_email,
@@ -175,6 +177,7 @@ def process_job_row(template_type, task_args_kwargs):
     send_fn.apply_async(
         *task_args_kwargs,
         queue=QueueNames.DATABASE,
+        MessageGroupId=message_group_id,
     )
 
 
@@ -353,6 +356,7 @@ def save_sms(
             provider_tasks.deliver_sms.apply_async(
                 [str(saved_notification.id)],
                 queue=QueueNames.SEND_SMS,
+                MessageGroupId=self.message_group_id,
             )
         else:
             extra = {
@@ -437,6 +441,7 @@ def save_email(self, service_id, notification_id, encoded_notification, sender_i
         provider_tasks.deliver_email.apply_async(
             [str(saved_notification.id)],
             queue=QueueNames.SEND_EMAIL,
+            MessageGroupId=self.message_group_id,
         )
 
         extra = {
@@ -494,7 +499,9 @@ def save_letter(
         )
 
         letters_pdf_tasks.get_pdf_for_templated_letter.apply_async(
-            [str(saved_notification.id)], queue=QueueNames.CREATE_LETTERS_PDF
+            [str(saved_notification.id)],
+            queue=QueueNames.CREATE_LETTERS_PDF,
+            MessageGroupId=self.message_group_id,
         )
 
         extra = {
@@ -529,7 +536,12 @@ def handle_exception(task, notification, notification_id, exc):
         # send to the retry queue.
         current_app.logger.exception("Retry: " + base_msg, extra, extra=extra)  # noqa
         try:
-            task.retry(queue=QueueNames.RETRY, exc=exc)
+            retry_kwargs = {
+                "queue": QueueNames.RETRY,
+                "exc": exc,
+                "MessageGroupId": task.message_group_id,
+            }
+            task.retry(**retry_kwargs)
         except task.MaxRetriesExceededError:
             current_app.logger.error("Max retry failed: " + base_msg, extra, extra=extra)  # noqa
 
@@ -580,7 +592,7 @@ def process_incomplete_job(job_id, shatter_batch_size=DEFAULT_SHATTER_JOB_ROWS_B
             get_id_task_args_kwargs_for_job_row(row, template, job, job.service, sender_id=sender_id)[1]
             for row in shatter_batch
         ]
-        _shatter_job_rows_with_subdivision(template.template_type, batch_args_kwargs)
+        _shatter_job_rows_with_subdivision(template.template_type, batch_args_kwargs, str(job.service_id))
 
     job_complete(job, resumed=True)
 
@@ -625,7 +637,11 @@ def _check_and_queue_returned_letter_callback_task(notification_id, service_id):
     # queue callback task only if the service_callback_api exists
     if service_callback_api := get_returned_letter_callback_api_for_service(service_id=service_id):
         returned_letter_data = create_returned_letter_callback_data(notification_id, service_id, service_callback_api)
-        send_returned_letter_to_service.apply_async([returned_letter_data], queue=QueueNames.CALLBACKS)
+        send_returned_letter_to_service.apply_async(
+            [returned_letter_data],
+            queue=QueueNames.CALLBACKS,
+            MessageGroupId=str(service_id),
+        )
 
 
 @notify_celery.task(bind=True, name="process-report-request")

app/job/rest.py

@@ -176,7 +176,12 @@ def create_job(service_id):
     sender_id = data.get("sender_id")
 
     if job.job_status == JOB_STATUS_PENDING:
-        process_job.apply_async([str(job.id)], {"sender_id": sender_id}, queue=QueueNames.JOBS)
+        process_job.apply_async(
+            [str(job.id)],
+            {"sender_id": sender_id},
+            queue=QueueNames.JOBS,
+            MessageGroupId=str(job.service_id),
+        )
 
     job_json = job_schema.dump(job)
     job_json["statistics"] = []

tests/app/celery/test_scheduled_tasks.py

@@ -73,7 +73,7 @@
     create_template,
     create_user,
 )
-from tests.conftest import set_config, set_config_values
+from tests.conftest import _with_message_group_id, set_config, set_config_values
 
 
 def test_should_call_delete_codes_on_delete_verify_codes_task(notify_db_session, mocker):
@@ -98,7 +98,11 @@ def test_should_update_scheduled_jobs_and_put_on_queue(mock_celery_task, sample_
 
     updated_job = dao_get_job_by_id(job.id)
     assert updated_job.job_status == "pending"
-    mocked.assert_called_with([str(job.id)], queue="job-tasks")
+    mocked.assert_called_with(
+        [str(job.id)],
+        queue="job-tasks",
+        MessageGroupId=str(job.service_id),
+    )
 
 
 def test_should_update_all_scheduled_jobs_and_put_on_queue(sample_template, mock_celery_task):
@@ -117,11 +121,12 @@ def test_should_update_all_scheduled_jobs_and_put_on_queue(sample_template, mock
     assert dao_get_job_by_id(job_2.id).job_status == "pending"
     assert dao_get_job_by_id(job_2.id).job_status == "pending"
 
+    service_id = str(sample_template.service_id)
     mocked.assert_has_calls(
         [
-            call([str(job_3.id)], queue="job-tasks"),
-            call([str(job_2.id)], queue="job-tasks"),
-            call([str(job_1.id)], queue="job-tasks"),
+            call([str(job_3.id)], queue="job-tasks", MessageGroupId=service_id),
+            call([str(job_2.id)], queue="job-tasks", MessageGroupId=service_id),
+            call([str(job_1.id)], queue="job-tasks", MessageGroupId=service_id),
         ]
     )
 
@@ -406,8 +411,11 @@ def test_check_job_status_task_sets_jobs_to_error(mock_celery_task, sample_templ
 
 
 def test_replay_created_notifications(sample_service, mock_celery_task):
-    email_delivery_queue = mock_celery_task(deliver_email)
-    sms_delivery_queue = mock_celery_task(deliver_sms)
+    with _with_message_group_id(deliver_email, str(sample_service.id)):
+        email_delivery_queue = mock_celery_task(deliver_email)
+
+    with _with_message_group_id(deliver_sms, str(sample_service.id)):
+        sms_delivery_queue = mock_celery_task(deliver_sms)
 
     sms_template = create_template(service=sample_service, template_type="sms")
     email_template = create_template(service=sample_service, template_type="email")
@@ -430,8 +438,12 @@ def test_replay_created_notifications(sample_service, mock_celery_task):
     create_notification(template=email_template, created_at=datetime.utcnow(), status="created")
 
     replay_created_notifications()
-    email_delivery_queue.assert_called_once_with([str(old_email.id)], queue="send-email-tasks")
-    sms_delivery_queue.assert_called_once_with([str(old_sms.id)], queue="send-sms-tasks")
+    email_delivery_queue.assert_called_once_with(
+        [str(old_email.id)], queue="send-email-tasks", MessageGroupId=str(sample_service.id)
+    )
+    sms_delivery_queue.assert_called_once_with(
+        [str(old_sms.id)], queue="send-sms-tasks", MessageGroupId=str(sample_service.id)
+    )
 
 
 def test_replay_created_notifications_get_pdf_for_templated_letter_tasks_for_letters_not_ready_to_send(
@@ -454,9 +466,10 @@ def test_replay_created_notifications_get_pdf_for_templated_letter_tasks_for_let
 
     replay_created_notifications()
 
+    service_id = str(sample_letter_template.service_id)
     calls = [
-        call([str(notification_1.id)], queue=QueueNames.CREATE_LETTERS_PDF),
-        call([str(notification_2.id)], queue=QueueNames.CREATE_LETTERS_PDF),
+        call([str(notification_1.id)], queue=QueueNames.CREATE_LETTERS_PDF, MessageGroupId=service_id),
+        call([str(notification_2.id)], queue=QueueNames.CREATE_LETTERS_PDF, MessageGroupId=service_id),
     ]
     mock_task.assert_has_calls(calls, any_order=True)
 
@@ -512,7 +525,10 @@ def test_check_if_letters_still_pending_virus_check_restarts_scan_for_stuck_lett
     mock_file_exists.assert_called_once_with("test-letters-scan", expected_filename)
 
     mock_celery.assert_called_once_with(
-        name=TaskNames.SCAN_FILE, kwargs={"filename": expected_filename}, queue=QueueNames.ANTIVIRUS
+        name=TaskNames.SCAN_FILE,
+        kwargs={"filename": expected_filename},
+        queue=QueueNames.ANTIVIRUS,
+        MessageGroupId=str(sample_letter_template.service_id),
     )
 
     assert mock_create_ticket.called is False
@@ -724,7 +740,12 @@ def test_check_for_missing_rows_in_completed_jobs(mocker, sample_email_template,
         )
     ]
     assert mock_save_email.mock_calls == [
-        mock.call((str(job.service_id), "some-uuid", "something_encoded"), {}, queue="database-tasks")
+        mock.call(
+            (str(job.service_id), "some-uuid", "something_encoded"),
+            {},
+            queue="database-tasks",
+            MessageGroupId=str(job.service_id),
+        )
     ]
 
 
@@ -765,7 +786,10 @@ def test_check_for_missing_rows_in_completed_jobs_uses_sender_id(
     ]
     assert mock_save_email.mock_calls == [
         mock.call(
-            (str(job.service_id), "some-uuid", "something_encoded"), {"sender_id": fake_uuid}, queue="database-tasks"
+            (str(job.service_id), "some-uuid", "something_encoded"),
+            {"sender_id": fake_uuid},
+            queue="database-tasks",
+            MessageGroupId=str(job.service_id),
         )
     ]