feat: Add MessageGroupId to job and save task queueing

- Pass MessageGroupId (service id) when calling shatter_job_rows.apply_async
  so batches use the same group id
- Add _service_id_from_job_row_args_kwargs helper and use it in
  _shatter_job_rows_with_subdivision and process_job_row
- Update tests to expect MessageGroupId on shatter, deliver, and retry
  calls; add tests for the new helper

Author: joybytes

app/celery/tasks.py

@@ -122,13 +122,16 @@ def process_job(job_id, sender_id=None, shatter_batch_size=DEFAULT_SHATTER_JOB_R
 
 
 def _shatter_job_rows_with_subdivision(template_type, args_kwargs_seq, top_level=True):
+    service_id = str(_service_id_from_job_row_args_kwargs(args_kwargs_seq[0])) if args_kwargs_seq else None
+
     try:
         shatter_job_rows.apply_async(
             (
                 template_type,
                 args_kwargs_seq,
             ),
             queue=QueueNames.JOBS,
+            MessageGroupId=service_id,
         )
     except BotoClientError as e:
         # this information is helpfully not preserved outside the message string of the exception, so
@@ -172,10 +175,9 @@ def process_job_row(template_type, task_args_kwargs):
         LETTER_TYPE: save_letter,
     }[template_type]
 
-    send_fn.apply_async(
-        *task_args_kwargs,
-        queue=QueueNames.DATABASE,
-    )
+    service_id = _service_id_from_job_row_args_kwargs(task_args_kwargs)
+    apply_kwargs = {"queue": QueueNames.DATABASE, "MessageGroupId": str(service_id) if service_id else None}
+    send_fn.apply_async(*task_args_kwargs, **apply_kwargs)
 
 
 def job_complete(job, resumed=False, start=None):
@@ -242,6 +244,12 @@ def get_id_task_args_kwargs_for_job_row(row, template, job, service, sender_id=N
     return notification_id, (task_args, task_kwargs)
 
 
+def _service_id_from_job_row_args_kwargs(task_args_kwargs):
+    """Extract service_id from (task_args, task_kwargs) produced by get_id_task_args_kwargs_for_job_row."""
+    task_args, _ = task_args_kwargs
+    return task_args[0]
+
+
 def __sending_limits_for_job_exceeded(service, job, job_id):
     try:
         check_service_over_daily_message_limit(
@@ -353,6 +361,7 @@ def save_sms(
             provider_tasks.deliver_sms.apply_async(
                 [str(saved_notification.id)],
                 queue=QueueNames.SEND_SMS,
+                MessageGroupId=str(service.id),
             )
         else:
             extra = {
@@ -437,6 +446,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=str(service.id),
         )
 
         extra = {
@@ -494,7 +504,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=str(service.id),
         )
 
         extra = {
@@ -529,7 +541,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
 
@@ -625,7 +642,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=str(job.service_id),
+        )
     ]
 
 
@@ -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=str(job.service_id),
         )
     ]