diff --git a/tcms_tenants/management/commands/remove_stale_tenant_attachments.py b/tcms_tenants/management/commands/remove_stale_tenant_attachments.py index a845c89..2136146 100644 --- a/tcms_tenants/management/commands/remove_stale_tenant_attachments.py +++ b/tcms_tenants/management/commands/remove_stale_tenant_attachments.py @@ -4,30 +4,24 @@ # https://www.gnu.org/licenses/agpl-3.0.html from django.core.files.storage import default_storage -from django.core.management.base import BaseCommand from django.utils import timezone from django_tenants.utils import get_tenant_model, tenant_context +from attachments.management.commands.delete_stale_attachments import ( + Command as DeleteStaleAttachmentsCommand, +) from attachments.models import Attachment from attachments.views import remove_file_from_disk -class Command(BaseCommand): +class Command(DeleteStaleAttachmentsCommand): help = ( "Remove attachments for which the related objects don't exist anymore! " "Works on all tenants!" ) def add_arguments(self, parser): - parser.add_argument( - "-y", - "--yes", - default="x", - action="store_const", - const="y", - dest="answer", - help="Automatically confirm deletion", - ) + super().add_arguments(parser) parser.add_argument( "--dry-run", action="store_true", @@ -75,6 +69,11 @@ def handle(self, *args, **kwargs): if kwargs["verbosity"]: output = self.stdout + if dry_run: + # the upstream command reports each match before prompting, so + # refusing to delete turns it into a dry run + kwargs["answer"] = "n" + for tenant in get_tenant_model().objects.all(): if output: output.write( @@ -83,28 +82,22 @@ def handle(self, *args, **kwargs): ) with tenant_context(tenant): - for attachment in Attachment.objects.all(): - if not attachment.object_id or attachment.content_object is None: - self.prompt_and_remove( - attachment, - f"Attachment `{attachment}' to non-existing " - f"`{attachment.content_type.model}' with PK " - f"`{attachment.object_id}'", - output, - dry_run, - answer, - ) - elif check_storage and not default_storage.exists( - attachment.attachment_file.name - ): - self.prompt_and_remove( - attachment, - f"Attachment `{attachment}' with missing file", - output, - dry_run, - answer, - delete_file=False, - ) + # upstream removes attachments whose related object is gone + super().handle(*args, **kwargs) + + if check_storage: + for attachment in Attachment.objects.all(): + if attachment.content_object and not ( + default_storage.exists(attachment.attachment_file.name) + ): + self.prompt_and_remove( + attachment, + f"Attachment `{attachment}' with missing file", + output, + dry_run, + answer, + delete_file=False, + ) if output: output.write( diff --git a/tcms_tenants/tests/test_remove_stale_tenant_attachments.py b/tcms_tenants/tests/test_remove_stale_tenant_attachments.py index 0cd342d..9b65cdd 100644 --- a/tcms_tenants/tests/test_remove_stale_tenant_attachments.py +++ b/tcms_tenants/tests/test_remove_stale_tenant_attachments.py @@ -76,17 +76,6 @@ def add_attachments(cls, parent_obj, index): stale.object_id = str(int(parent_obj.pk) + cls.pk_offset) stale.save() - stale2 = Attachment.objects.create( - content_type=content_type, - object_id=parent_obj.pk, - attachment_file=SimpleUploadedFile( - f"stale2{index}.txt", b"attachment content" - ), - creator=cls.tester, - ) - stale2.object_id = "" - stale2.save() - def test_removes_only_attachments_which_are_attached_to_missing_objects(self): tenants = tenants_with_schema() self.assertGreaterEqual(len(tenants), 1) @@ -94,13 +83,14 @@ def test_removes_only_attachments_which_are_attached_to_missing_objects(self): for tenant in tenants: with tenant_context(tenant): self.assertTrue( - Attachment.objects.exclude(object_id="") - .annotate(object_id_as_int=Cast("object_id", IntegerField())) + Attachment.objects.annotate( + object_id_as_int=Cast("object_id", IntegerField()) + ) .filter(object_id_as_int__gt=self.pk_offset) .exists() ) self.assertEqual( - Attachment.objects.count(), self.attachments_per_model * 3 * 3 + Attachment.objects.count(), self.attachments_per_model * 2 * 3 ) out = StringIO() @@ -119,8 +109,9 @@ def test_removes_only_attachments_which_are_attached_to_missing_objects(self): Attachment.objects.count(), self.attachments_per_model * 3 ) self.assertFalse( - Attachment.objects.exclude(object_id="") - .annotate(object_id_as_int=Cast("object_id", IntegerField())) + Attachment.objects.annotate( + object_id_as_int=Cast("object_id", IntegerField()) + ) .filter(object_id_as_int__gt=self.pk_offset) .exists() ) @@ -149,6 +140,39 @@ def test_removes_only_attachments_which_are_attached_to_missing_objects(self): output, ) + def test_dry_run_removes_nothing(self): + tenants = tenants_with_schema() + self.assertGreaterEqual(len(tenants), 1) + + out = StringIO() + call_command( + "remove_stale_tenant_attachments", + answer="y", + dry_run=True, + verbosity=1, + stdout=out, + ) + + output = out.getvalue() + + for tenant in tenants: + with tenant_context(tenant): + self.assertEqual( + Attachment.objects.count(), self.attachments_per_model * 2 * 3 + ) + self.assertTrue( + Attachment.objects.annotate( + object_id_as_int=Cast("object_id", IntegerField()) + ) + .filter(object_id_as_int__gt=self.pk_offset) + .exists() + ) + + self.assertIn( + f"for tenant '{tenant.schema_name}' === dry run: True ===", + output, + ) + def test_removes_attachments_whose_file_is_missing(self): with tenant_context(self.tenant): test_case = TestCaseFactory()