feat: propagate MessageGroupId through job processing tasks

- Make process_job and shatter_job_rows bind=True and pass message_group_id
  through to child tasks via getattr(self, message_group_id, None)
- Add message_group_id to process_job_row and save_* apply_async calls
- process_incomplete_job passes service_id as message_group_id
- Update tests: expect None when tasks run directly; add test for propagation

Author: joybytes

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": getattr(task, "message_group_id", None),
+            }
+            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")

tests/app/celery/test_scheduled_tasks.py

@@ -724,7 +724,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=None,
+        )
     ]
 
 
@@ -765,7 +770,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=None,
         )
     ]