paperless-ngx/src/documents/plugins/helpers.py

import enum
from typing import TYPE_CHECKING
from typing import Literal
from typing import Self
from typing import TypeAlias
from typing import TypedDict

from asgiref.sync import async_to_sync
from channels.layers import get_channel_layer

if TYPE_CHECKING:
    from channels_redis.pubsub import RedisPubSubChannelLayer


class ProgressStatusOptions(enum.StrEnum):
    STARTED = "STARTED"
    WORKING = "WORKING"
    SUCCESS = "SUCCESS"
    FAILED = "FAILED"


class PermissionsData(TypedDict, total=False):
    """Permission fields included in status messages for access control."""

    owner_id: int | None
    users_can_view: list[int]
    groups_can_view: list[int]


class ProgressUpdateData(TypedDict):
    filename: str | None
    task_id: str | None
    current_progress: int
    max_progress: int
    status: str
    message: str
    document_id: int | None
    owner_id: int | None
    users_can_view: list[int]
    groups_can_view: list[int]


class StatusUpdatePayload(TypedDict):
    type: Literal["status_update"]
    data: ProgressUpdateData


class DocumentsDeletedData(TypedDict):
    documents: list[int]


class DocumentsDeletedPayload(TypedDict):
    type: Literal["documents_deleted"]
    data: DocumentsDeletedData


class DocumentUpdatedData(TypedDict):
    document_id: int
    modified: str
    owner_id: int | None
    users_can_view: list[int]
    groups_can_view: list[int]


class DocumentUpdatedPayload(TypedDict):
    type: Literal["document_updated"]
    data: DocumentUpdatedData


WebsocketPayload: TypeAlias = (
    StatusUpdatePayload | DocumentsDeletedPayload | DocumentUpdatedPayload
)


class BaseStatusManager:
    """
    Handles sending of progress information via the channel layer, with proper management
    of the open/close of the layer to ensure messages go out and everything is cleaned up
    """

    def __init__(self) -> None:
        self._channel: RedisPubSubChannelLayer | None = None

    def __enter__(self) -> Self:
        self.open()
        return self

    def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None:
        self.close()

    def open(self) -> None:
        """
        If not already opened, gets the default channel layer
        opened and ready to send messages
        """
        if self._channel is None:
            self._channel = get_channel_layer()

    def close(self) -> None:
        """
        If it was opened, flushes the channel layer
        """
        if self._channel is not None:
            async_to_sync(self._channel.flush)
            self._channel = None

    def send(self, payload: WebsocketPayload) -> None:
        # Ensure the layer is open
        self.open()

        # Just for IDEs
        if TYPE_CHECKING:
            assert self._channel is not None

        # Construct and send the update
        async_to_sync(self._channel.group_send)("status_updates", payload)


class ProgressManager(BaseStatusManager):
    def __init__(self, filename: str | None = None, task_id: str | None = None) -> None:
        super().__init__()
        self.filename = filename
        self.task_id = task_id

    def send_progress(
        self,
        status: ProgressStatusOptions,
        message: str,
        current_progress: int,
        max_progress: int,
        *,
        document_id: int | None = None,
        owner_id: int | None = None,
        users_can_view: list[int] | None = None,
        groups_can_view: list[int] | None = None,
    ) -> None:
        data: ProgressUpdateData = {
            "filename": self.filename,
            "task_id": self.task_id,
            "current_progress": current_progress,
            "max_progress": max_progress,
            "status": status,
            "message": message,
            "document_id": document_id,
            "owner_id": owner_id,
            "users_can_view": users_can_view or [],
            "groups_can_view": groups_can_view or [],
        }
        payload: StatusUpdatePayload = {"type": "status_update", "data": data}
        self.send(payload)


class DocumentsStatusManager(BaseStatusManager):
    def send_documents_deleted(self, documents: list[int]) -> None:
        payload: DocumentsDeletedPayload = {
            "type": "documents_deleted",
            "data": {
                "documents": documents,
            },
        }
        self.send(payload)

    def send_document_updated(
        self,
        *,
        document_id: int,
        modified: str,
        owner_id: int | None = None,
        users_can_view: list[int] | None = None,
        groups_can_view: list[int] | None = None,
    ) -> None:
        payload: DocumentUpdatedPayload = {
            "type": "document_updated",
            "data": {
                "document_id": document_id,
                "modified": modified,
                "owner_id": owner_id,
                "users_can_view": users_can_view or [],
                "groups_can_view": groups_can_view or [],
            },
        }
        self.send(payload)