feat: management command to delete a product

Author: StefanFl

backend/application/core/management/commands/delete_product.py

@@ -0,0 +1,83 @@
+import logging
+from argparse import ArgumentParser
+from typing import Any
+
+from django.core.management.base import BaseCommand, CommandError
+from django.db import transaction
+from django.db.models import Model
+from django.db.models.signals import post_delete
+
+from application.core.models import Observation, Product
+from application.core.signals import observation_post_delete
+from application.licenses.models import License_Component
+from application.licenses.signals import license_component_post_delete
+
+logger = logging.getLogger("secobserve.core")
+
+
+class Command(BaseCommand):
+
+    help = "Delete a product with many observations in batches"
+
+    def add_arguments(self, parser: ArgumentParser) -> None:
+        parser.add_argument("product_name", help="Exact name of the product")
+        parser.add_argument("--batch-size", type=int, default=1000, help="Number of objects deleted per transaction")
+        parser.add_argument("--no-input", action="store_true", help="Do not ask for confirmation")
+
+    def handle(self, *args: Any, **options: Any) -> None:
+        product_name = options["product_name"]
+        batch_size = options["batch_size"]
+        if batch_size < 1:
+            raise CommandError("--batch-size must be at least 1")
+
+        product = Product.objects.filter(name=product_name).first()
+        if not product:
+            raise CommandError(f"Product {product_name} not found")
+        if product.is_product_group:
+            raise CommandError(f"{product_name} is a product group, delete the products of the group first")
+
+        observation_count = Observation.objects.filter(product=product).count()
+        license_component_count = License_Component.objects.filter(product=product).count()
+        logger.info(
+            "Product %s has %s observations and %s license components",
+            product_name,
+            observation_count,
+            license_component_count,
+        )
+
+        if not options["no_input"]:
+            answer = input(f"Delete product {product_name}? [y/N] ")
+            if answer.strip().lower() != "y":
+                logger.info("... aborted")
+                return
+
+        # The receivers are not needed when the whole product is deleted and would slow down the deletion
+        post_delete.disconnect(observation_post_delete, sender=Observation)
+        post_delete.disconnect(license_component_post_delete, sender=License_Component)
+        try:
+            self._delete_in_batches(Observation, product, batch_size, observation_count)
+            self._delete_in_batches(License_Component, product, batch_size, license_component_count)
+
+            with transaction.atomic():
+                product.delete()
+        finally:
+            post_delete.connect(observation_post_delete, sender=Observation)
+            post_delete.connect(license_component_post_delete, sender=License_Component)
+
+        logger.info("... product %s deleted", product_name)
+
+    def _delete_in_batches(self, model: type[Model], product: Product, batch_size: int, total: int) -> None:
+        name = model.__name__
+        deleted = 0
+        while True:
+            ids = list(
+                model.objects.filter(product=product)  # type: ignore[attr-defined]
+                .order_by("id")
+                .values_list("id", flat=True)[:batch_size]
+            )
+            if not ids:
+                break
+            with transaction.atomic():
+                model.objects.filter(id__in=ids).delete()  # type: ignore[attr-defined]
+            deleted += len(ids)
+            logger.info("... %s / %s %s deleted", deleted, total, name)

backend/unittests/core/management/commands/test_delete_product.py

@@ -0,0 +1,127 @@
+from datetime import datetime, timezone
+from unittest.mock import patch
+
+from django.core.management import call_command
+from django.core.management.base import CommandError
+from django.test import TestCase
+from django.utils import timezone as django_timezone
+
+from application.access_control.models import User
+from application.core.models import (
+    Branch,
+    Evidence,
+    Observation,
+    Observation_Log,
+    Product,
+)
+from application.import_observations.models import Parser
+from application.licenses.models import License_Component
+
+
+class TestDeleteProductCommand(TestCase):
+    """DB-backed tests for the `delete_product` management command."""
+
+    def setUp(self) -> None:
+        self.user = User.objects.create(username="delete-product@example.com")
+        self.parser = Parser.objects.create(name="parser_1")
+        self.product = Product.objects.create(name="product_1")
+        self.other_product = Product.objects.create(name="product_2")
+        self.branch = Branch.objects.create(product=self.product, name="main")
+
+        for i in range(3):
+            observation = self._create_observation(self.product, f"observation_{i}")
+            Observation_Log.objects.create(observation=observation, user=self.user, comment="comment")
+            Evidence.objects.create(observation=observation, name="evidence", evidence="evidence")
+        self._create_license_component(self.product)
+
+        self.other_observation = self._create_observation(self.other_product, "other_observation")
+        self.other_license_component = self._create_license_component(self.other_product)
+
+    def _create_observation(self, product: Product, title: str) -> Observation:
+        return Observation.objects.create(
+            title=title,
+            product=product,
+            parser=self.parser,
+            numerical_severity=1,
+            import_last_seen=django_timezone.now(),
+        )
+
+    def _create_license_component(self, product: Product) -> License_Component:
+        return License_Component.objects.create(
+            product=product,
+            component_name="component",
+            numerical_evaluation_result=3,
+        )
+
+    def test_product_deleted(self) -> None:
+        call_command("delete_product", "product_1", "--no-input", "--batch-size", "2")
+
+        self.assertFalse(Product.objects.filter(name="product_1").exists())
+        self.assertFalse(Observation.objects.filter(product=self.product).exists())
+        self.assertFalse(Observation_Log.objects.filter(observation__product=self.product).exists())
+        self.assertFalse(Evidence.objects.filter(observation__product=self.product).exists())
+        self.assertFalse(License_Component.objects.filter(product=self.product).exists())
+        self.assertFalse(Branch.objects.filter(product=self.product).exists())
+
+        self.assertTrue(Product.objects.filter(name="product_2").exists())
+        self.assertTrue(Observation.objects.filter(pk=self.other_observation.pk).exists())
+        self.assertTrue(License_Component.objects.filter(pk=self.other_license_component.pk).exists())
+
+    def test_product_not_found(self) -> None:
+        with self.assertRaisesMessage(CommandError, "Product unknown not found"):
+            call_command("delete_product", "unknown", "--no-input")
+
+    def test_product_group(self) -> None:
+        product_group = Product.objects.create(name="product_group", is_product_group=True)
+        self.product.product_group = product_group
+        self.product.save()
+
+        with self.assertRaisesMessage(CommandError, "product_group is a product group"):
+            call_command("delete_product", "product_group", "--no-input")
+
+        self.assertTrue(Product.objects.filter(name="product_group").exists())
+        self.assertEqual(3, Observation.objects.filter(product=self.product).count())
+
+    def test_invalid_batch_size(self) -> None:
+        with self.assertRaisesMessage(CommandError, "--batch-size must be at least 1"):
+            call_command("delete_product", "product_1", "--no-input", "--batch-size", "0")
+
+    @patch("builtins.input", return_value="n")
+    def test_confirmation_declined(self, mock_input) -> None:
+        call_command("delete_product", "product_1")
+
+        mock_input.assert_called_once()
+        self.assertTrue(Product.objects.filter(name="product_1").exists())
+        self.assertEqual(3, Observation.objects.filter(product=self.product).count())
+
+    @patch("builtins.input", return_value="y")
+    def test_confirmation_accepted(self, mock_input) -> None:
+        call_command("delete_product", "product_1")
+
+        mock_input.assert_called_once()
+        self.assertFalse(Product.objects.filter(name="product_1").exists())
+
+    @patch("application.core.signals.push_deleted_observation_to_issue_tracker")
+    def test_no_issue_tracker_push(self, mock_push) -> None:
+        self.product.issue_tracker_active = True
+        self.product.save()
+        Observation.objects.filter(product=self.product).update(issue_tracker_issue_id="123")
+
+        with self.captureOnCommitCallbacks(execute=True):
+            call_command("delete_product", "product_1", "--no-input")
+
+        mock_push.assert_not_called()
+
+    @patch("application.core.signals.push_deleted_observation_to_issue_tracker")
+    def test_receivers_reconnected(self, mock_push) -> None:
+        call_command("delete_product", "product_1", "--no-input")
+
+        with self.captureOnCommitCallbacks(execute=True):
+            self.other_observation.delete()
+        mock_push.assert_called_once()
+
+        old_timestamp = datetime(2020, 1, 1, tzinfo=timezone.utc)
+        Product.objects.filter(pk=self.other_product.pk).update(last_license_change=old_timestamp)
+        License_Component.objects.get(pk=self.other_license_component.pk).delete()
+        self.other_product.refresh_from_db()
+        self.assertGreater(self.other_product.last_license_change, old_timestamp)