feat(category): faster removal

docs/changes.rst

@@ -18,6 +18,7 @@ Weblate 5.17
 * :ref:`addon-weblate.cdn.cdnjs` validates parsed locations.
 * Removed unintended API endpoints for translation memory.
 * Improved API access control for pending tasks.
+* Faster category removal.
 
 .. rubric:: Compatibility
 

weblate/trans/models/__init__.py

@@ -29,6 +29,7 @@
 from weblate.trans.models.unit import Unit
 from weblate.trans.models.variant import Variant
 from weblate.trans.models.workflow import WorkflowSetting
+from weblate.trans.removal import get_current_removal_batch
 from weblate.trans.signals import user_pre_delete
 from weblate.utils.decorators import disable_for_loaddata
 from weblate.utils.files import remove_tree
@@ -84,6 +85,10 @@ def project_post_delete(sender, instance: Project, **kwargs) -> None:
 
 @receiver(pre_delete, sender=Component)
 def component_pre_delete(sender, instance: Component, **kwargs) -> None:
+    batch = instance.removal_batch or get_current_removal_batch()
+    if batch is not None:
+        batch.collect_stats(instance.stats.get_update_objects())
+        return
     # Collect list of stats to update, this can't be done after removal
     instance.stats.collect_update_objects()
 
@@ -96,8 +101,9 @@ def component_post_delete(sender, instance: Component, **kwargs) -> None:
     - delete component directory
     - update stats, this is accompanied by component_pre_delete
     """
-    # Update stats
-    transaction.on_commit(instance.stats.update_parents)
+    batch = instance.removal_batch or get_current_removal_batch()
+    if batch is None:
+        transaction.on_commit(instance.stats.update_parents)
     instance.stats.delete()
 
     # Do not delete linked components
@@ -200,6 +206,10 @@ def auto_component_list(sender, instance, **kwargs) -> None:
 @receiver(post_delete, sender=Component)
 @disable_for_loaddata
 def post_delete_linked(sender, instance, **kwargs) -> None:
+    batch = instance.removal_batch or get_current_removal_batch()
+    if batch is not None:
+        batch.collect_linked_component(instance.linked_component_id)
+        return
     # When removing project, the linked component might be already deleted now
     with suppress(Component.DoesNotExist):
         if instance.linked_component:

weblate/trans/models/component.py

@@ -150,6 +150,7 @@
     from weblate.lang.models import Plural
     from weblate.trans.models import Project
     from weblate.trans.models.unit import UnitAttributesDict
+    from weblate.trans.removal import RemovalBatch
     from weblate.vcs.base import Repository
 
 NEW_LANG_CHOICES = (
@@ -920,6 +921,7 @@
         self.translations_count: int | None = None
         self.translations_progress = 0
         self.acting_user: User | None = None
+        self.removal_batch: RemovalBatch | None = None
         self.batch_checks = False
         self.batched_checks: set[str] = set()
         self.needs_variants_update = False
@@ -1233,7 +1235,7 @@
         # Calculate progress for translations
         if progress is None:
             self.translations_progress += 1
             progress = 100 * self.translations_progress // self.translations_count
         # Store task state
         current_task.update_state(
             state="PROGRESS", meta={"progress": progress, "component": self.pk}
@@ -1275,7 +1277,7 @@
         if "source_translation" in self.__dict__:
             return self.__dict__["source_translation"]
         try:
             result = self.translation_set.get(language_id=self.source_language_id)
         except ObjectDoesNotExist:
             return None
         self.__dict__["source_translation"] = result
@@ -1309,7 +1311,7 @@
                     )
             except IntegrityError:
                 try:
                     return self.translation_set.get(language_id=self.source_language_id)
                 except self.translation_set.model.DoesNotExist:
                     pass
                 raise
@@ -1339,7 +1341,7 @@
 
     def _process_new_source(self, source: Unit, *, save: bool = True) -> Change:
         # Avoid fetching empty list of checks from the database
         source.all_checks = []
         source.source_updated = True
         change = source.generate_change(
             self.acting_user,
@@ -2821,7 +2823,7 @@
         if self.has_template():
             # Avoid parsing if template is invalid
             try:
                 self.template_store.check_valid()
             except (ValueError, FileParseError) as exc:
                 raise InvalidTemplateError(info=str(exc)) from exc
         self._template_check_done = True
@@ -2885,7 +2887,7 @@
             if (
                 self.file_format == "po"
                 and self.new_base.endswith(".pot")
                 and os.path.exists(self.get_new_base_filename())
             ):
                 # Process pot file as source to include additional metadata
                 matches = [self.new_base, *matches]
@@ -2954,7 +2956,7 @@
                     )
                     self.handle_parse_error(error.__cause__, filename=self.template)
                     self.update_import_alerts()
                     raise error.__cause__ from error  # pylint: disable=raising-non-exception
                 was_change |= bool(translation.reason)
                 translations[translation.id] = translation
                 languages[lang.code] = translation
@@ -3287,7 +3289,7 @@
         filename = self.get_new_base_filename()
         template = self.has_template()
         return self.file_format_cls.is_valid_base_for_new(
             filename,
             template,
             errors,
             fast=fast,
@@ -3312,7 +3314,7 @@
             raise ValidationError({"new_base": message, "new_lang": message})
         filename = self.get_new_base_filename()
         # File is present, but does not exist
         if not os.path.exists(filename):
             raise ValidationError({"new_base": gettext("File does not exist.")})
         # File is present, but it is not valid
         if errors:
@@ -3393,7 +3395,7 @@
                 raise ValidationError({"template": gettext("File does not exist.")})
 
             try:
                 self.template_store.check_valid()
             except (FileParseError, ValueError) as error:
                 msg = gettext("Could not parse translation base file: %s") % str(error)
                 raise ValidationError({"template": msg}) from error

weblate/trans/removal.py

@@ -0,0 +1,60 @@
+# Copyright © Michal Čihař <michal@weblate.org>
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator
+
+    from weblate.utils.stats import BaseStats
+
+CURRENT_REMOVAL_BATCH: ContextVar[RemovalBatch | None] = ContextVar(
+    "current_removal_batch", default=None
+)
+
+
+class RemovalBatch:
+    def __init__(self) -> None:
+        self.removed_component_ids: set[int] = set()
+        self.stats_to_update: dict[str, BaseStats] = {}
+        self.linked_components_to_refresh: set[int] = set()
+
+    def mark_component(self, component_id: int) -> None:
+        self.removed_component_ids.add(component_id)
+
+    def collect_stats(self, stats_objects: Iterable[BaseStats]) -> None:
+        for stats in stats_objects:
+            self.stats_to_update[stats.cache_key] = stats
+
+    def collect_linked_component(self, component_id: int | None) -> None:
+        if component_id is not None:
+            self.linked_components_to_refresh.add(component_id)
+
+    def flush(self) -> None:
+        from weblate.trans.models import Component
+
+        for stats in self.stats_to_update.values():
+            stats.update_stats()
+
+        for component in Component.objects.filter(
+            pk__in=self.linked_components_to_refresh
+        ):
+            component.update_alerts()
+
+
+def get_current_removal_batch() -> RemovalBatch | None:
+    return CURRENT_REMOVAL_BATCH.get()
+
+
+@contextmanager
+def removal_batch_context(batch: RemovalBatch) -> Iterator[None]:
+    token = CURRENT_REMOVAL_BATCH.set(batch)
+    try:
+        yield
+    finally:
+        CURRENT_REMOVAL_BATCH.reset(token)

weblate/trans/tasks.py

@@ -44,6 +44,7 @@
     Translation,
 )
 from weblate.trans.models.unit import fill_in_source_translation
+from weblate.trans.removal import RemovalBatch, removal_batch_context
 from weblate.utils.celery import app
 from weblate.utils.data import data_dir
 from weblate.utils.errors import report_error
@@ -439,6 +440,15 @@ def component_removal(pk: int, uid: int) -> None:
         component = Component.objects.get(pk=pk)
     except Component.DoesNotExist:
         return
+
+    _component_removal(component, user)
+
+
+def _component_removal(
+    component: Component, user: User, batch: RemovalBatch | None = None
+) -> None:
+    if batch is not None:
+        component.removal_batch = batch
     with component.repository.lock:
         component.acting_user = user
         component.project.change_set.create(
@@ -452,22 +462,44 @@ def component_removal(pk: int, uid: int) -> None:
             components = component.project.component_set.filter(
                 allow_translation_propagation=True
             ).exclude(pk=component.pk)
-            for component in components.iterator():
-                component.schedule_update_checks()
+            if batch is not None:
+                components = components.exclude(pk__in=batch.removed_component_ids)
+            for current in components.iterator():
+                current.schedule_update_checks()
 
 
-@app.task(trail=False)
-@transaction.atomic
-def category_removal(pk: int, uid: int) -> None:
-    user = User.objects.get(pk=uid)
-    try:
-        category = Category.objects.get(pk=pk)
-    except Category.DoesNotExist:
-        return
+def _collect_removal_targets(category: Category, batch: RemovalBatch) -> None:
+    linked_frontier: set[int] = set()
+    component_ids = category.component_set.values_list("id", flat=True).iterator(
+        chunk_size=1000
+    )
+    for component_id in component_ids:
+        batch.mark_component(component_id)
+        linked_frontier.add(component_id)
+
+    while linked_frontier:
+        children = Component.objects.filter(
+            linked_component_id__in=linked_frontier
+        ).values_list("id", flat=True)
+        next_frontier = set()
+        for component_id in children.iterator(chunk_size=1000):
+            if component_id in batch.removed_component_ids:
+                continue
+            batch.mark_component(component_id)
+            next_frontier.add(component_id)
+        linked_frontier = next_frontier
+
     for child in category.category_set.all():
-        category_removal(child.pk, uid)
-    for component_id in category.component_set.values_list("id", flat=True):
-        component_removal(component_id, uid)
+        _collect_removal_targets(child, batch)
+
+
+def _category_removal(
+    category: Category, user: User, batch: RemovalBatch | None = None
+) -> None:
+    for child in category.category_set.all():
+        _category_removal(child, user, batch)
+    for component in category.component_set.iterator():
+        _component_removal(component, user, batch)
     category.project.change_set.create(
         action=ActionEvents.REMOVE_CATEGORY,
         target=category.slug,
@@ -477,6 +509,21 @@ def category_removal(pk: int, uid: int) -> None:
     category.delete()
 
 
+@app.task(trail=False)
+@transaction.atomic
+def category_removal(pk: int, uid: int) -> None:
+    user = User.objects.get(pk=uid)
+    try:
+        category = Category.objects.get(pk=pk)
+    except Category.DoesNotExist:
+        return
+    batch = RemovalBatch()
+    _collect_removal_targets(category, batch)
+    with removal_batch_context(batch):
+        _category_removal(category, user, batch)
+    transaction.on_commit(batch.flush)
+
+
 @app.task(
     trail=False,
     autoretry_for=(IntegrityError,),

weblate/trans/tests/test_categories.py

@@ -5,12 +5,18 @@
 """Test for categories."""
 
 import os
+from contextlib import ExitStack
+from unittest.mock import call, patch
 
 from django.urls import reverse
 
 from weblate.lang.models import get_default_lang
+from weblate.trans.actions import ActionEvents
 from weblate.trans.models import Category, Component, Project
+from weblate.trans.removal import RemovalBatch
+from weblate.trans.tasks import category_removal
 from weblate.trans.tests.test_views import ViewTestCase
+from weblate.utils.stats import GlobalStats
 
 
 class CategoriesTest(ViewTestCase):
@@ -316,3 +322,131 @@ def test_move_category_linked_repo(self) -> None:
 
         component.refresh_from_db()
         self.assertEqual(component.repo, "weblate://other/test-rename/test")
+
+    def test_category_removal_batches_linked_alert_updates(self) -> None:
+        category = Category.objects.create(
+            project=self.project, name="Category test", slug="testcat"
+        )
+        first = self.create_link_existing(
+            name="Linked A", slug="linked-a", category=category
+        )
+        second = self.create_link_existing(
+            name="Linked B", slug="linked-b", category=category
+        )
+
+        with patch.object(Component, "update_alerts", autospec=True) as update_alerts:
+            category_removal(category.pk, self.user.pk)
+
+        self.assertFalse(
+            Component.objects.filter(pk__in=[first.pk, second.pk]).exists()
+        )
+        update_alerts.assert_called_once_with(self.component)
+
+    def test_category_removal_skips_propagation_for_removed_components(self) -> None:
+        category = Category.objects.create(
+            project=self.project, name="Category test", slug="testcat"
+        )
+        first = self.create_po(
+            project=self.project,
+            name="Category A",
+            slug="category-a",
+            category=category,
+        )
+        second = self.create_po(
+            project=self.project,
+            name="Category B",
+            slug="category-b",
+            category=category,
+        )
+        first.allow_translation_propagation = True
+        first.save(update_fields=["allow_translation_propagation"])
+        second.allow_translation_propagation = True
+        second.save(update_fields=["allow_translation_propagation"])
+
+        with patch.object(
+            Component, "schedule_update_checks", autospec=True
+        ) as schedule:
+            category_removal(category.pk, self.user.pk)
+
+        self.assertEqual(
+            [call(self.component), call(self.component)],
+            schedule.call_args_list,
+        )
+
+    def test_category_removal_skips_propagation_for_cascaded_linked_components(
+        self,
+    ) -> None:
+        category = Category.objects.create(
+            project=self.project, name="Category test", slug="testcat"
+        )
+        self.component.category = category
+        self.component.save()
+
+        linked = self.create_link_existing(name="Linked A", slug="linked-a")
+        linked.allow_translation_propagation = True
+        linked.save(update_fields=["allow_translation_propagation"])
+
+        other = self.create_po(
+            project=self.project,
+            name="Category A",
+            slug="category-a",
+            category=category,
+        )
+        other.allow_translation_propagation = True
+        other.save(update_fields=["allow_translation_propagation"])
+
+        with patch.object(
+            Component, "schedule_update_checks", autospec=True
+        ) as schedule:
+            category_removal(category.pk, self.user.pk)
+
+        schedule.assert_not_called()
+
+    def test_category_removal_batches_parent_stats_updates(self) -> None:
+        category = Category.objects.create(
+            project=self.project, name="Category test", slug="testcat"
+        )
+        for pos in range(2):
+            self.create_po(
+                project=self.project,
+                name=f"Category {pos}",
+                slug=f"category-{pos}",
+                category=category,
+            )
+
+        collected: list[set[str]] = []
+        executed: list[str] = []
+        original_flush = RemovalBatch.flush
+
+        def record_flush(batch_self) -> None:
+            collected.append(set(batch_self.stats_to_update))
+            with ExitStack() as stack:
+                for stats in batch_self.stats_to_update.values():
+                    stack.enter_context(
+                        patch.object(
+                            stats,
+                            "update_stats",
+                            side_effect=lambda stats=stats: executed.append(
+                                stats.cache_key
+                            ),
+                        )
+                    )
+                original_flush(batch_self)
+
+        with patch.object(
+            RemovalBatch, "flush", autospec=True, side_effect=record_flush
+        ):
+            category_removal(category.pk, self.user.pk)
+
+        self.assertEqual(1, len(collected))
+        self.assertEqual(
+            {category.stats.cache_key, GlobalStats().cache_key},
+            collected[0],
+        )
+        self.assertEqual(collected[0], set(executed))
+        self.assertEqual(
+            2,
+            self.project.change_set.filter(
+                action=ActionEvents.REMOVE_COMPONENT
+            ).count(),
+        )