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