Ping via ws after bulk operations complete

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")
 

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()

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))