From 736983b7e1f8c9f49578fe8157992c969c0835ba Mon Sep 17 00:00:00 2001 From: shamoon <4887959+shamoon@users.noreply.github.com> Date: Tue, 3 Nov 2026 09:02:52 -0800 Subject: [PATCH] Ping via ws after bulk operations complete --- src/documents/signals/handlers.py | 28 ++++++ src/documents/tests/test_bulk_operations.py | 94 +++++++++++++++++++++ src/paperless/consumers.py | 15 ++++ 3 files changed, 137 insertions(+) diff --git a/src/documents/signals/handlers.py b/src/documents/signals/handlers.py index 1cd520646e..805b3886b9 100644 --- a/src/documents/signals/handlers.py +++ b/src/documents/signals/handlers.py @@ -8,6 +8,7 @@ from pathlib import Path from typing import TYPE_CHECKING from typing import Any +from asgiref.sync import async_to_sync from celery import shared_task from celery.signals import before_task_publish from celery.signals import task_failure @@ -16,6 +17,7 @@ from celery.signals import task_prerun from celery.signals import task_revoked from celery.signals import worker_process_init from celery.signals import worker_process_shutdown +from channels.layers import get_channel_layer from django.conf import settings from django.contrib.auth.models import Group from django.contrib.auth.models import User @@ -1242,6 +1244,29 @@ def task_prerun_handler(sender=None, task_id=None, task=None, **kwargs) -> None: logger.exception("Setting PaperlessTask started failed") +def _send_bulk_operation_status(task_id: str) -> None: + """Notify the initiator, task records still happen.""" + try: + task = PaperlessTask.objects.filter(task_id=task_id).first() + if task is None or task.owner_id is None or "operation" not in task.input_data: + return + async_to_sync(get_channel_layer().group_send)( + f"bulk_operations_{task.owner_id}", + { + "type": "bulk_operation_status", + "data": { + "task_id": task.task_id, + "status": task.status, + "operation": task.input_data["operation"], + "error": (task.result_data or {}).get("error_message", ""), + }, + }, + ) + except Exception: + # Just a notification failure + logger.exception("Sending bulk operation status failed") + + @task_postrun.connect def task_postrun_handler( sender=None, @@ -1299,6 +1324,7 @@ def task_postrun_handler( changed_fields.append("status") task_instance.save(update_fields=changed_fields) + _send_bulk_operation_status(task_id) except Exception: # pragma: no cover logger.exception("Updating PaperlessTask failed") @@ -1364,6 +1390,7 @@ def task_failure_handler( date_started - date_created ).total_seconds() task_qs.update(**update_fields) + _send_bulk_operation_status(task_id) except Exception: # pragma: no cover logger.exception("Updating PaperlessTask on failure failed") @@ -1407,6 +1434,7 @@ def task_revoked_handler( "reason": (reason + " Check worker logs for details."), }, ) + _send_bulk_operation_status(task_id) except Exception: # pragma: no cover logger.exception("Updating PaperlessTask on revocation failed") diff --git a/src/documents/tests/test_bulk_operations.py b/src/documents/tests/test_bulk_operations.py index 6a4190372c..d9cff036bf 100644 --- a/src/documents/tests/test_bulk_operations.py +++ b/src/documents/tests/test_bulk_operations.py @@ -182,3 +182,97 @@ def test_large_filtered_selection_is_one_task(bulk_setup): assert response.status_code == 200 enqueue.assert_called_once() assert len(enqueue.call_args.kwargs["kwargs"]["document_ids"]) == 10_000 + + +@pytest.mark.django_db +@pytest.mark.parametrize("failed", [False, True]) +def test_tracked_outcome_notifies_only_initiator(bulk_setup, failed): + user, _, doc, _ = bulk_setup + task_id = str(uuid.uuid4()) + before_task_publish_handler( + headers={ + "task": "documents.bulk_operations.execute_bulk_operation", + "id": task_id, + }, + body=( + (), + {"operation": "delete", "document_ids": [doc.pk], "user_id": user.pk}, + {}, + ), + ) + task = PaperlessTask.objects.get(task_id=task_id) + assert task.owner_id == user.pk + assert task.task_type == PaperlessTask.TaskType.BULK_DELETE + with mock.patch("documents.signals.handlers.async_to_sync") as send: + if failed: + task_failure_handler( + task_id=task_id, + exception=ValueError("Could not update"), + ) + else: + task_postrun_handler( + task_id=task_id, + state="SUCCESS", + retval={"document_count": 1}, + ) + group, event = send.return_value.call_args.args + assert group == f"bulk_operations_{user.pk}" + assert event["data"]["status"] == ("failure" if failed else "success") + assert event["data"]["task_id"] == task_id + + +@pytest.mark.django_db +def test_bulk_status_is_available_without_task_permission_and_is_private(bulk_setup): + _, client, _, _ = bulk_setup + owner = UserFactory() + other = UserFactory() + owned = PaperlessTask.objects.create( + task_id=str(uuid.uuid4()), + task_type="bulk_update", + owner=owner, + trigger_source="web_ui", + input_data={"operation": "set_permissions"}, + ) + foreign = PaperlessTask.objects.create( + task_id=str(uuid.uuid4()), + task_type="bulk_update", + owner=other, + trigger_source="web_ui", + input_data={"operation": "set_permissions"}, + ) + client.force_authenticate(owner) + response = client.get( + "/api/tasks/bulk_status/", + {"task_id": [owned.task_id, foreign.task_id]}, + ) + assert response.status_code == 200 + assert [item["task_id"] for item in response.data] == [owned.task_id] + + +@pytest.mark.anyio +@pytest.mark.parametrize("anyio_backend", ["asyncio"]) +async def test_completion_socket_is_private(anyio_backend): + from types import SimpleNamespace + + from channels.layers import get_channel_layer + from channels.testing import WebsocketCommunicator + + from paperless.consumers import StatusConsumer + + own = WebsocketCommunicator(StatusConsumer.as_asgi(), "/ws/status/") + other = WebsocketCommunicator(StatusConsumer.as_asgi(), "/ws/status/") + own.scope["user"] = SimpleNamespace(id=1, is_authenticated=True) + other.scope["user"] = SimpleNamespace(id=2, is_authenticated=True) + assert (await own.connect())[0] + assert (await other.connect())[0] + try: + event = { + "type": "bulk_operation_status", + "data": {"task_id": "example", "status": "success"}, + } + await get_channel_layer().group_send("bulk_operations_1", event) + assert await own.receive_json_from() == event + assert await other.receive_nothing(timeout=0.02) + finally: + await own.disconnect() + await other.disconnect() diff --git a/src/paperless/consumers.py b/src/paperless/consumers.py index 5f9a8891ed..63c78c5b2b 100644 --- a/src/paperless/consumers.py +++ b/src/paperless/consumers.py @@ -45,12 +45,21 @@ class StatusConsumer(AsyncWebsocketConsumer): await self.close() return await self.channel_layer.group_add("status_updates", self.channel_name) + await self.channel_layer.group_add( + f"bulk_operations_{self.scope['user'].id}", + self.channel_name, + ) await self.accept() self._start_heartbeat() async def disconnect(self, code: int) -> None: await self._stop_heartbeat() await self.channel_layer.group_discard("status_updates", self.channel_name) + if self._authenticated(): + await self.channel_layer.group_discard( + f"bulk_operations_{self.scope['user'].id}", + self.channel_name, + ) def _start_heartbeat(self) -> None: self.heartbeat_task = asyncio.create_task(self._heartbeat_loop()) @@ -84,3 +93,9 @@ class StatusConsumer(AsyncWebsocketConsumer): await self.close() elif await self._can_view(event["data"]): await self.send(json.dumps(event)) + + async def bulk_operation_status(self, event: dict) -> None: + if not self._authenticated(): + await self.close() + else: + await self.send(json.dumps(event))