leocrm/app/core/restore_registry.py

"""Entity restore registry — declarative configuration for undo/restore.

Each registered entity type declares:
- model_class: SQLAlchemy model to load
- restore_permission: permission string required to restore
- excluded_fields: fields never restored from snapshot (id, tenant_id, timestamps, etc.)
- special_handler: optional async callable for entity-specific restore logic

No dynamic ORM loading, no blind snapshot writes — only explicitly registered
entity types can be restored, and only through their declared configuration.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Any, Awaitable, Callable

from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)

# Default fields excluded from restore for every entity type
_DEFAULT_EXCLUDED = frozenset({
    "id",
    "tenant_id",
    "created_at",
    "updated_at",
    "deleted_at",
    "search_tsv",
    "embedding",
})

# Type alias for special restore handler
SpecialRestoreHandler = Callable[
    [AsyncSession, Any, str, Any, dict[str, Any]],
    Awaitable[dict[str, Any]],
]


@dataclass
class RestoreConfig:
    """Configuration for restoring a specific entity type."""

    entity_type: str
    model_class: type
    restore_permission: str
    excluded_fields: frozenset[str] = field(default_factory=frozenset)
    special_handler: SpecialRestoreHandler | None = None

    @property
    def all_excluded_fields(self) -> frozenset[str]:
        """Merge default excluded fields with entity-specific ones."""
        return _DEFAULT_EXCLUDED | self.excluded_fields


class RestoreRegistry:
    """Singleton registry mapping entity_type → RestoreConfig."""

    _instance: RestoreRegistry | None = None

    def __new__(cls) -> RestoreRegistry:
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._configs: dict[str, RestoreConfig] = {}
        return cls._instance

    def register(self, config: RestoreConfig) -> None:
        """Register a RestoreConfig for an entity type."""
        if config.entity_type in self._configs:
            logger.warning("Overwriting restore config for entity_type: %s", config.entity_type)
        self._configs[config.entity_type] = config
        logger.debug("Registered restore config for: %s", config.entity_type)

    def get(self, entity_type: str) -> RestoreConfig | None:
        """Get RestoreConfig for entity_type, or None if not registered."""
        return self._configs.get(entity_type)

    def is_registered(self, entity_type: str) -> bool:
        """Check if entity_type is registered for restore."""
        return entity_type in self._configs

    def list_registered(self) -> list[str]:
        """Return all registered entity types."""
        return sorted(self._configs.keys())

    def _reset_for_testing(self) -> None:
        """Clear all registrations — for unit tests only."""
        self._configs.clear()


def get_restore_registry() -> RestoreRegistry:
    """Return the global RestoreRegistry singleton."""
    return RestoreRegistry()


def reset_restore_registry_for_testing() -> RestoreRegistry:
    """Return a fresh singleton — for unit tests only."""
    reg = get_restore_registry()
    reg._reset_for_testing()
    return reg


# ─── Default entity registrations ───


def register_default_entities() -> None:
    """Register all built-in entity types for restore.

    Called during app startup. Plugin entities should register themselves
    in their on_activate() lifecycle hook.
    """
    from app.models.contact import Contact

    reg = get_restore_registry()

    # Contact (covers both 'person' and 'company' types — same model)
    reg.register(RestoreConfig(
        entity_type="contact",
        model_class=Contact,
        restore_permission="contacts:write",
        excluded_fields=frozenset({
            "search_tsv",
            "embedding",
            "default_person_id",
            "admin_contactperson_id",
        }),
    ))

    # Task plugin
    try:
        from app.plugins.builtins.tasks.models import Task
        reg.register(RestoreConfig(
            entity_type="task",
            model_class=Task,
            restore_permission="tasks:write",
            excluded_fields=frozenset({
                "created_by",
                "assigned_to",
                "contact_id",
            }),
        ))
    except ImportError:
        logger.debug("Tasks plugin model not available for restore registration")

    # Calendar plugin — CalendarEntry
    try:
        from app.plugins.builtins.calendar.models import CalendarEntry
        reg.register(RestoreConfig(
            entity_type="calendar_entry",
            model_class=CalendarEntry,
            restore_permission="calendar:write",
            excluded_fields=frozenset({
                "calendar_id",
                "created_by",
                "assigned_to",
                "source_mail_id",
            }),
        ))
    except ImportError:
        logger.debug("Calendar plugin model not available for restore registration")

    # DMS plugin — File metadata
    try:
        from app.plugins.builtins.dms.models import File as DmsFile
        reg.register(RestoreConfig(
            entity_type="dms_file",
            model_class=DmsFile,
            restore_permission="dms:write",
            excluded_fields=frozenset({
                "storage_path",
                "content_hash",
                "size_bytes",
                "uploaded_by",
                "folder_id",
            }),
        ))
    except ImportError:
        logger.debug("DMS plugin model not available for restore registration")

    # Mail plugin — special handler for IMAP semantics
    try:
        from app.plugins.builtins.mail.models import Mail
        reg.register(RestoreConfig(
            entity_type="mail",
            model_class=Mail,
            restore_permission="mail:write",
            excluded_fields=frozenset({
                "message_id",
                "rfc822_size",
                "raw_path",
                "account_id",
                "folder_id",
            }),
            special_handler=_mail_restore_handler,
        ))
    except ImportError:
        logger.debug("Mail plugin model not available for restore registration")


async def _mail_restore_handler(
    db: AsyncSession,
    entity: Any,
    action: str,
    snapshot: dict[str, Any],
    context: dict[str, Any],
) -> dict[str, Any]:
    """Special restore handler for Mail entities.

    Mail restore has IMAP semantics:
    - delete: move back from trash to original folder (if folder still exists)
    - update: revert metadata fields
    - create: soft-delete (undo send only works for drafts)

    Server errors must not produce false local status.
    """
    import uuid
    from datetime import datetime, timezone
    from sqlalchemy import select

    user_id = context.get("user_id")
    tenant_id = context.get("tenant_id")

    if action == "delete":
        # Un-delete: clear deleted_at, restore original folder_id if available
        if entity is None:
            raise ValueError("Mail entity not found for restore")
        entity.deleted_at = None
        if user_id:
            entity.updated_by = user_id if hasattr(entity, "updated_by") else None
        # Restore original folder from snapshot if available
        original_folder_id = snapshot.get("folder_id")
        if original_folder_id and hasattr(entity, "folder_id"):
            try:
                folder_uuid = uuid.UUID(str(original_folder_id))
                # Verify folder still exists and is not deleted
                from app.plugins.builtins.mail.models import MailFolder
                folder_q = select(MailFolder).where(
                    MailFolder.id == folder_uuid,
                    MailFolder.tenant_id == tenant_id,
                    MailFolder.deleted_at.is_(None),
                )
                folder_result = await db.execute(folder_q)
                folder = folder_result.scalar_one_or_none()
                if folder:
                    entity.folder_id = folder_uuid
                else:
                    logger.warning(
                        "Original mail folder %s no longer exists, "
                        "restoring mail without folder assignment",
                        original_folder_id,
                    )
            except (ValueError, Exception) as e:
                logger.warning("Failed to restore mail folder: %s", e)

        await db.flush()
        return {"id": str(entity.id), "restored": True, "entity_type": "mail"}

    elif action == "update":
        if entity is None:
            raise ValueError("Mail entity not found for restore")
        # Revert metadata fields from snapshot_before
        excluded = _DEFAULT_EXCLUDED | {
            "message_id", "rfc822_size", "raw_path", "account_id", "folder_id",
        }
        for key, value in snapshot.items():
            if hasattr(entity, key) and key not in excluded:
                setattr(entity, key, value)
        await db.flush()
        return {"id": str(entity.id), "restored": True, "entity_type": "mail"}

    elif action == "create":
        # Undo creation: soft-delete (only meaningful for drafts)
        if entity is None:
            raise ValueError("Mail entity not found for restore")
        entity.deleted_at = datetime.now(timezone.utc)
        await db.flush()
        return {"id": str(entity.id), "restored": True, "entity_type": "mail", "note": "soft-deleted (undo create)"}

    raise ValueError(f"Unsupported action for mail restore: {action}")