Prevent deletion of file based models if object is in use

wagtail_to_ion/admin.py

@@ -4,13 +4,30 @@
 from wagtail_to_ion.models import get_ion_media_rendition_model
 
 
-class AbstractIonImageAdmin(admin.ModelAdmin):
+class ProtectInUseModelAdmin(admin.ModelAdmin):
+    protect_objects_in_use: bool = True
+
+    def get_deleted_objects(self, objs, request):
+        if self.protect_objects_in_use:
+            protected = []
+
+            for obj in objs:
+                usage = obj.get_usage()
+                if usage:
+                    protected += list(usage)
+            if protected:
+                return [], {}, set(), protected
+
+        return super().get_deleted_objects(objs, request)
+
+
+class AbstractIonImageAdmin(ProtectInUseModelAdmin):
     list_display = ('title', 'collection', 'include_in_archive')
     list_filter = ('collection', 'include_in_archive')
     search_fields = ('title',)
 
 
-class AbstractIonDocumentAdmin(admin.ModelAdmin):
+class AbstractIonDocumentAdmin(ProtectInUseModelAdmin):
     list_display = ('title', 'collection', 'include_in_archive')
     list_filter = ('collection', 'include_in_archive')
     search_fields = ('title',)
@@ -21,7 +38,7 @@ class IonMediaRenditionInlineAdmin(admin.TabularInline):
     extra = 0
 
 
-class AbstractIonMediaAdmin(admin.ModelAdmin):
+class AbstractIonMediaAdmin(ProtectInUseModelAdmin):
     list_display = ('title', 'collection', 'include_in_archive')
     list_filter = ('collection', 'include_in_archive')
     search_fields = ('title',)

wagtail_to_ion/blocks.py

@@ -0,0 +1,17 @@
+from django.utils.functional import cached_property
+from wagtailmedia.blocks import AbstractMediaChooserBlock
+
+
+class IonMediaBlock(AbstractMediaChooserBlock):
+    @cached_property
+    def target_model(self):
+        from wagtailmedia.models import get_media_model
+        return get_media_model()
+
+    @cached_property
+    def widget(self):
+        from wagtailmedia.widgets import AdminMediaChooser
+        return AdminMediaChooser
+
+    def render_basic(self, value, context=None):
+        raise NotImplementedError('You need to implement %s.render_basic' % self.__class__.__name__)

wagtail_to_ion/models/file_based_models.py

@@ -3,15 +3,19 @@
 import os
 
 from django.db import models
-from django.utils.functional import cached_property
+from django.db.models import ProtectedError
+from django.db.models.signals import post_delete, pre_delete
+from django.dispatch import receiver
 from django.utils.translation import ugettext_lazy as _
 
 from magic import from_buffer as magic_from_buffer
+from wagtail.documents.blocks import DocumentChooserBlock
 from wagtail.documents.models import AbstractDocument
+from wagtail.images.blocks import ImageChooserBlock
 from wagtail.images.models import AbstractImage, AbstractRendition
-from wagtailmedia.blocks import AbstractMediaChooserBlock
 from wagtailmedia.models import AbstractMedia
 
+from wagtail_to_ion.blocks import IonMediaBlock
 from wagtail_to_ion.conf import settings
 from wagtail_to_ion.models import get_ion_media_rendition_model
 from wagtail_to_ion.tasks import generate_media_rendition, get_audio_metadata
@@ -32,6 +36,7 @@ class AbstractIonDocument(AbstractDocument):
         'tags',
         'include_in_archive',
     )
+    check_usage_block_types = (DocumentChooserBlock,)
 
     class Meta:
         abstract = True
@@ -52,6 +57,10 @@ def save(self, *args, **kwargs):
         super().save(*args, **kwargs)
         os.chmod(self.file.path, 0o644)
 
+    def get_usage(self):
+        from wagtail_to_ion.utils import get_object_block_usage
+        return super().get_usage().union(get_object_block_usage(self, block_types=self.check_usage_block_types))
+
 
 class AbstractIonImage(AbstractImage):
     checksum = models.CharField(max_length=255)
@@ -70,6 +79,7 @@ class AbstractIonImage(AbstractImage):
         'focal_point_height',
         'include_in_archive',
     )
+    check_usage_block_types = (ImageChooserBlock,)
 
     class Meta:
         abstract = True
@@ -90,6 +100,10 @@ def save(self, *args, **kwargs):
         super().save(*args, **kwargs)
         os.chmod(self.file.path, 0o644)
 
+    def get_usage(self):
+        from wagtail_to_ion.utils import get_object_block_usage
+        return super().get_usage().union(get_object_block_usage(self, block_types=self.check_usage_block_types))
+
     @property
     def archive_rendition(self):
         result = self.get_rendition(self.rendition_type)
@@ -158,6 +172,7 @@ class AbstractIonMedia(AbstractMedia):
         'collection',
         'tags',
     )
+    check_usage_block_types = (IonMediaBlock,)
 
     class Meta:
         abstract = True
@@ -203,19 +218,9 @@ def save(self, *args, **kwargs):
         if needs_transcode:
             self.create_renditions()
 
-    def delete(self, *args, **kwargs):
-        try:
-            self.file.delete()
-        except ValueError:
-            pass
-        try:
-            self.thumbnail.delete()
-        except ValueError:
-            pass
-        # delete one by one to make sure files are deleted
-        for rendition in self.renditions.all():
-            rendition.delete()
-        super().delete(*args, **kwargs)
+    def get_usage(self):
+        from wagtail_to_ion.utils import get_object_block_usage
+        return super().get_usage().union(get_object_block_usage(self, block_types=self.check_usage_block_types))
 
     def set_media_metadata(self):
         self.file.open()
@@ -297,28 +302,27 @@ class Meta:
     def __str__(self):
         return "IonMediaRendition {} for {}".format(self.name, self.media_item)
 
-    def delete(self, *args, **kwargs):
+
+@receiver(pre_delete)
+def prevent_deletion_if_in_use(sender, instance, **kwargs):
+    if isinstance(instance, (AbstractIonDocument, AbstractIonImage, AbstractIonMedia)):
+        usage = instance.get_usage()
+        if usage:
+            model_name = instance.__class__.__name__
+            raise ProtectedError(
+                f"Cannot delete instance of model '{model_name}' because it is referenced in stream field blocks",
+                usage,
+            )
+
+
+@receiver(post_delete)
+def remove_media_files(sender, instance, **kwargs):
+    if isinstance(instance, (AbstractIonMedia, AbstractIonMediaRendition)):
         try:
-            self.file.delete()
+            instance.file.delete(save=False)
         except ValueError:
             pass
         try:
-            self.thumbnail.delete()
+            instance.thumbnail.delete(save=False)
         except ValueError:
             pass
-        super().delete(*args, **kwargs)
-
-
-class IonMediaBlock(AbstractMediaChooserBlock):
-    @cached_property
-    def target_model(self):
-        from wagtailmedia.models import get_media_model
-        return get_media_model()
-
-    @cached_property
-    def widget(self):
-        from wagtailmedia.widgets import AdminMediaChooser
-        return AdminMediaChooser
-
-    def render_basic(self, value, context=None):
-        raise NotImplementedError('You need to implement %s.render_basic' % self.__class__.__name__)

wagtail_to_ion/utils.py

@@ -1,6 +1,11 @@
-from django.db.models import Q
+import functools
+from typing import Any, Dict, Generator, List, NamedTuple, Tuple, Type, Union
 
-from wagtail.core.models import Collection, Page, PageViewRestriction
+from django.db.models import Q, Model
+
+from wagtail.core.blocks import Block, BoundBlock, StreamValue, StructBlock, StructValue
+from wagtail.core.fields import StreamField
+from wagtail.core.models import Collection, Page, PageViewRestriction, get_page_models
 from wagtail.images import get_image_model
 from wagtail.documents import get_document_model
 
@@ -61,7 +66,7 @@ def get_collection_for_page(page):
 # TODO: might be obsolete once https://github.com/wagtail/wagtail/pull/6300 has been merged
 def visible_tree_by_user(root, user):
     collection = IonCollection.objects.get(live=True, slug=get_collection_for_page(root))
-    
+
     if collection.view_restrictions.exists():
         restrictions = collection.view_restrictions.filter(
             restriction_type=PageViewRestriction.GROUPS,
@@ -115,3 +120,110 @@ def isoDate(d):
         return d.replace(microsecond=0, tzinfo=None).strftime("%Y-%m-%dT%H:%M:%SZ")
     else:
         return 'None'
+
+
+def get_object_block_usage(obj, block_types: Union[Type[Block], Tuple[Type[Block]]]):
+    """
+    Returns a queryset of pages that contain a block linked to a particular object.
+
+    Works like `wagtail.admin.models.get_object_usage` but inspects all pages
+    which might contain a block of the specified type(s).
+
+    To avoid inspecting all pages the following optimizations are applied:
+      1. operate only on page models with a StreamField containing the specified block type(s)
+      2. filter pages by expected JSON string in StreamField column(s)
+    """
+    block_usage = get_page_models_using_blocks(block_types=block_types)
+    page_ptr_ids = set()
+
+    for page_model in block_usage.keys():
+        stream_field_filter_q = Q()
+        stream_fields = set()
+
+        # create a (postgres specific) filter for every block; look for:
+        # - `"value": <pk>` string if block is in a StreamValue
+        # - `"<field_name>": <pk>` if block is in a StructValue
+        for block in block_usage[page_model]:
+            filter_att = block.block_name if block.in_struct else 'value'
+            stream_field_filter_q |= Q(**{f'{block.stream_field_name}__regex': rf'"{filter_att}":\s*{obj.pk}\M'})
+            stream_fields.add(block.stream_field_name)
+
+        for page_with_obj_pk_in_blocks in page_model.objects.filter(stream_field_filter_q):
+            for field_name in stream_fields:
+                for bound_block, _ in get_stream_value_blocks(getattr(page_with_obj_pk_in_blocks, field_name)):
+                    if isinstance(bound_block.block, block_types) and bound_block.value == obj:
+                        page_ptr_ids.add(page_with_obj_pk_in_blocks.page_ptr_id)
+
+    return Page.objects.filter(pk__in=page_ptr_ids)
+
+
+class StreamFieldBlockInfo(NamedTuple):
+    block_name: str
+    block_type: Block
+    in_struct: bool
+
+
+def get_stream_field_blocks(stream_field) -> Generator[StreamFieldBlockInfo, None, None]:
+    """Generates an un-nested list of blocks of a `StreamField`."""
+    def unnest_blocks(blocks: dict, in_struct: bool = False):
+        for block in blocks.items():
+            yield StreamFieldBlockInfo(block[0], block[1], in_struct)
+            if hasattr(block[1], 'child_blocks'):
+                yield from unnest_blocks(block[1].child_blocks, in_struct=isinstance(block[1], StructBlock))
+
+    return unnest_blocks(stream_field.stream_block.child_blocks)
+
+
+class StreamValueBlockInfo(NamedTuple):
+    bound_block: BoundBlock
+    json_value_field_name: str
+
+
+def get_stream_value_blocks(stream_value: StreamValue) -> Generator[StreamValueBlockInfo, None, None]:
+    """Generates an un-nested list of blocks of a `StreamValue`."""
+    def unnest_blocks(value: Union[StreamValue, StructValue]):
+        if isinstance(value, StreamValue):
+            for stream_block in value:
+                assert isinstance(stream_block, BoundBlock)
+                if isinstance(stream_block.value, StructValue):
+                    yield from unnest_blocks(stream_block.value)
+                else:
+                    yield StreamValueBlockInfo(stream_block, 'value')
+        elif isinstance(value, StructValue):
+            for block_name, bound_block in value.bound_blocks.items():
+                assert isinstance(bound_block, BoundBlock)
+                if isinstance(bound_block.value, StreamValue):
+                    yield from unnest_blocks(bound_block.value)
+                else:
+                    yield StreamValueBlockInfo(bound_block, block_name)
+        else:
+            raise RuntimeError(f'Unexpected type: {type(value)}')
+
+    return unnest_blocks(stream_value)
+
+
+class ModelStreamFieldBlockInfo(NamedTuple):
+    stream_field_name: str
+    block_name: str
+    block_type: Block
+    in_struct: bool
+
+
+@functools.lru_cache()
+def get_page_models_using_blocks(
+    block_types: Union[Type[Block], Tuple[Type[Block]]],
+) -> Dict[Type[Model], List[ModelStreamFieldBlockInfo]]:
+    """Returns information about all page models which use the specified block type(s) in a StreamField."""
+    models_with_block: Dict[Type[Model], List[ModelStreamFieldBlockInfo]] = {}
+
+    for page_model in get_page_models():
+        for stream_field in [field for field in page_model._meta.get_fields() if isinstance(field, StreamField)]:
+            for stream_field_block in get_stream_field_blocks(stream_field):
+                if isinstance(stream_field_block.block_type, block_types):
+                    if page_model not in models_with_block:
+                        models_with_block[page_model] = []
+                    models_with_block[page_model].append(
+                        ModelStreamFieldBlockInfo(stream_field.attname, *stream_field_block),
+                    )
+
+    return models_with_block

wagtail_to_ion/views/wagtail_override/__init__.py

@@ -1 +1,2 @@
 from .pages import ion_unpublish, ion_publish_with_children
+from .file_based_models import document_safe_delete, image_safe_delete, media_safe_delete

wagtail_to_ion/views/wagtail_override/file_based_models.py

@@ -0,0 +1,48 @@
+from django.shortcuts import get_object_or_404, redirect
+from django.utils.translation import gettext_lazy
+
+from wagtail.admin import messages
+from wagtail.admin.auth import PermissionPolicyChecker
+from wagtail.documents import get_document_model
+from wagtail.documents.permissions import permission_policy as document_permission_policy
+from wagtail.documents.views.documents import delete as document_delete_view
+from wagtail.images import get_image_model
+from wagtail.images.permissions import permission_policy as image_permission_policy
+from wagtail.images.views.images import delete as image_delete_view
+from wagtailmedia.models import get_media_model
+from wagtailmedia.permissions import permission_policy as media_permission_policy
+from wagtailmedia.views.media import delete as media_delete_view
+
+
+document_permission_checker = PermissionPolicyChecker(document_permission_policy)
+image_permission_checker = PermissionPolicyChecker(image_permission_policy)
+media_permission_checker = PermissionPolicyChecker(media_permission_policy)
+
+MESSAGE = gettext_lazy('Cannot delete object if used by any page')
+
+
+@image_permission_checker.require('delete')
+def image_safe_delete(request, image_id):
+    image = get_object_or_404(get_image_model(), pk=image_id)
+    if image.get_usage().exists():
+        messages.error(request, MESSAGE)
+        return redirect('wagtailimages:edit', image_id=image_id)
+    return image_delete_view(request, image_id)
+
+
+@document_permission_checker.require('delete')
+def document_safe_delete(request, document_id):
+    document = get_object_or_404(get_document_model(), pk=document_id)
+    if document.get_usage().exists():
+        messages.error(request, MESSAGE)
+        return redirect('wagtaildocs:edit', document_id=document_id)
+    return document_delete_view(request, document_id)
+
+
+@media_permission_checker.require('delete')
+def media_safe_delete(request, media_id):
+    media = get_object_or_404(get_media_model(), pk=media_id)
+    if media.get_usage().exists():
+        messages.error(request, MESSAGE)
+        return redirect('wagtailmedia:edit', media_id)
+    return media_delete_view(request, media_id)

wagtail_to_ion/wagtail_hooks/__init__.py

@@ -1,3 +1,4 @@
 from .show_only_user_images import show_only_user_images
 from .show_only_user_documents import show_only_user_documents
 from .dual_range_slider import global_admin_css, global_admin_js
+from .wagtail_admin_urls import register_admin_urls