Use celery chord for rotate tasks

This commit is contained in:
shamoon
2024-03-21 21:22:01 -07:00
parent cb4e300c77
commit f800632524
2 changed files with 13 additions and 7 deletions

View File

@@ -4,6 +4,7 @@ import logging
import os import os
from typing import Optional from typing import Optional
from celery import chord
from django.conf import settings from django.conf import settings
from django.db.models import Q from django.db.models import Q
@@ -167,6 +168,7 @@ def rotate(doc_ids: list[int], degrees: int):
affected_docs = [] affected_docs = []
import pikepdf import pikepdf
rotate_tasks = []
for doc in qs: for doc in qs:
try: try:
with pikepdf.open(doc.source_path, allow_overwriting_input=True) as pdf: with pikepdf.open(doc.source_path, allow_overwriting_input=True) as pdf:
@@ -175,8 +177,10 @@ def rotate(doc_ids: list[int], degrees: int):
pdf.save() pdf.save()
doc.checksum = hashlib.md5(doc.source_path.read_bytes()).hexdigest() doc.checksum = hashlib.md5(doc.source_path.read_bytes()).hexdigest()
doc.save() doc.save()
update_document_archive_file.delay( rotate_tasks.append(
document_id=doc.id, update_document_archive_file.s(
document_id=doc.id,
),
) )
logger.info( logger.info(
f"Rotated document {doc.id} by {degrees} degrees", f"Rotated document {doc.id} by {degrees} degrees",
@@ -186,7 +190,8 @@ def rotate(doc_ids: list[int], degrees: int):
logger.exception(f"Error rotating document {doc.id}: {e}") logger.exception(f"Error rotating document {doc.id}: {e}")
if len(affected_docs) > 0: if len(affected_docs) > 0:
bulk_update_documents.delay(document_ids=affected_docs) bulk_update_task = bulk_update_documents.s(document_ids=affected_docs)
chord(header=rotate_tasks, body=bulk_update_task).delay()
return "OK" return "OK"

View File

@@ -452,9 +452,10 @@ class TestPDFActions(DirectoriesMixin, TestCase):
mock_consume_file.assert_not_called() mock_consume_file.assert_not_called()
@mock.patch("documents.tasks.bulk_update_documents.delay") @mock.patch("documents.tasks.bulk_update_documents.s")
@mock.patch("documents.tasks.update_document_archive_file.delay") @mock.patch("documents.tasks.update_document_archive_file.s")
def test_rotate(self, mock_update_document, mock_update_documents): @mock.patch("celery.chord.delay")
def test_rotate(self, mock_chord, mock_update_document, mock_update_documents):
""" """
GIVEN: GIVEN:
- Existing documents - Existing documents
@@ -467,7 +468,7 @@ class TestPDFActions(DirectoriesMixin, TestCase):
result = bulk_edit.rotate(doc_ids, 90) result = bulk_edit.rotate(doc_ids, 90)
self.assertEqual(mock_update_document.call_count, 2) self.assertEqual(mock_update_document.call_count, 2)
mock_update_documents.assert_called_once() mock_update_documents.assert_called_once()
mock_chord.assert_called_once()
self.assertEqual(result, "OK") self.assertEqual(result, "OK")
@mock.patch("documents.tasks.bulk_update_documents.delay") @mock.patch("documents.tasks.bulk_update_documents.delay")