ref(providers): Refactor provider deletion functions (#7351)

Co-authored-by: Víctor Fernández Poyatos <victor@prowler.com>

api/CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to the **Prowler API** are documented in this file.
 
 ---
 
+## [v1.5.2] (Prowler v5.4.2)
+
+### Changed
+- Refactored deletion logic and implemented retry mechanism for deletion tasks [(#7349)](https://github.com/prowler-cloud/prowler/pull/7349).
+
+---
+
 ## [v1.5.1] (Prowler v5.4.1)
 
 ### Fixed

api/src/backend/api/db_utils.py

@@ -6,6 +6,7 @@ from datetime import datetime, timedelta, timezone
 from django.conf import settings
 from django.contrib.auth.models import BaseUserManager
 from django.db import connection, models, transaction
+from django_celery_beat.models import PeriodicTask
 from psycopg2 import connect as psycopg2_connect
 from psycopg2.extensions import AsIs, new_type, register_adapter, register_type
 from rest_framework_json_api.serializers import ValidationError
@@ -105,11 +106,12 @@ def generate_random_token(length: int = 14, symbols: str | None = None) -> str:
     return "".join(secrets.choice(symbols or _symbols) for _ in range(length))
 
 
-def batch_delete(queryset, batch_size=5000):
+def batch_delete(tenant_id, queryset, batch_size=5000):
     """
     Deletes objects in batches and returns the total number of deletions and a summary.
 
     Args:
+        tenant_id (str): Tenant ID the queryset belongs to.
         queryset (QuerySet): The queryset of objects to delete.
         batch_size (int): The number of objects to delete in each batch.
 
@@ -120,15 +122,16 @@ def batch_delete(queryset, batch_size=5000):
     deletion_summary = {}
 
     while True:
-        # Get a batch of IDs to delete
-        batch_ids = set(
-            queryset.values_list("id", flat=True).order_by("id")[:batch_size]
-        )
-        if not batch_ids:
-            # No more objects to delete
-            break
+        with rls_transaction(tenant_id, POSTGRES_TENANT_VAR):
+            # Get a batch of IDs to delete
+            batch_ids = set(
+                queryset.values_list("id", flat=True).order_by("id")[:batch_size]
+            )
+            if not batch_ids:
+                # No more objects to delete
+                break
 
-        deleted_count, deleted_info = queryset.filter(id__in=batch_ids).delete()
+            deleted_count, deleted_info = queryset.filter(id__in=batch_ids).delete()
 
         total_deleted += deleted_count
         for model_label, count in deleted_info.items():
@@ -137,6 +140,18 @@ def batch_delete(queryset, batch_size=5000):
     return total_deleted, deletion_summary
 
 
+def delete_related_daily_task(provider_id: str):
+    """
+    Deletes the periodic task associated with a specific provider.
+
+    Args:
+        provider_id (str): The unique identifier for the provider
+                           whose related periodic task should be deleted.
+    """
+    task_name = f"scan-perform-scheduled-{provider_id}"
+    PeriodicTask.objects.filter(name=task_name).delete()
+
+
 # Postgres Enums
 
 

api/src/backend/api/signals.py

@@ -1,12 +1,12 @@
 from celery import states
 from celery.signals import before_task_publish
+from config.celery import celery_app
 from django.db.models.signals import post_delete
 from django.dispatch import receiver
-from django_celery_beat.models import PeriodicTask
 from django_celery_results.backends.database import DatabaseBackend
 
+from api.db_utils import delete_related_daily_task
 from api.models import Provider
-from config.celery import celery_app
 
 
 def create_task_result_on_publish(sender=None, headers=None, **kwargs):  # noqa: F841
@@ -31,5 +31,4 @@ before_task_publish.connect(
 @receiver(post_delete, sender=Provider)
 def delete_provider_scan_task(sender, instance, **kwargs):  # noqa: F841
     # Delete the associated periodic task when the provider is deleted
-    task_name = f"scan-perform-scheduled-{instance.id}"
-    PeriodicTask.objects.filter(name=task_name).delete()
+    delete_related_daily_task(instance.id)

api/src/backend/api/specs/v1.yaml

@@ -1,7 +1,7 @@
 openapi: 3.0.3
 info:
   title: Prowler API
-  version: 1.5.1
+  version: 1.5.2
   description: |-
     Prowler API specification.
 

api/src/backend/api/tests/test_db_utils.py

@@ -131,9 +131,10 @@ class TestBatchDelete:
         return provider_count
 
     @pytest.mark.django_db
-    def test_batch_delete(self, create_test_providers):
+    def test_batch_delete(self, tenants_fixture, create_test_providers):
+        tenant_id = str(tenants_fixture[0].id)
         _, summary = batch_delete(
-            Provider.objects.all(), batch_size=create_test_providers // 2
+            tenant_id, Provider.objects.all(), batch_size=create_test_providers // 2
         )
         assert Provider.objects.all().count() == 0
         assert summary == {"api.Provider": create_test_providers}

api/src/backend/api/v1/views.py

@@ -55,6 +55,7 @@ from tasks.tasks import (
 
 from api.base_views import BaseRLSViewSet, BaseTenantViewset, BaseUserViewset
 from api.db_router import MainRouter
+from api.db_utils import delete_related_daily_task
 from api.filters import (
     ComplianceOverviewFilter,
     FindingFilter,
@@ -240,7 +241,7 @@ class SchemaView(SpectacularAPIView):
 
     def get(self, request, *args, **kwargs):
         spectacular_settings.TITLE = "Prowler API"
-        spectacular_settings.VERSION = "1.5.1"
+        spectacular_settings.VERSION = "1.5.2"
         spectacular_settings.DESCRIPTION = (
             "Prowler API specification.\n\nThis file is auto-generated."
         )
@@ -1077,6 +1078,7 @@ class ProviderViewSet(BaseRLSViewSet):
         provider = get_object_or_404(Provider, pk=pk)
         provider.is_deleted = True
         provider.save()
+        delete_related_daily_task(str(provider.id))
 
         with transaction.atomic():
             task = delete_provider_task.delay(

api/src/backend/config/celery.py

@@ -50,9 +50,9 @@ class RLSTask(Task):
 
         tenant_id = kwargs.get("tenant_id")
         with rls_transaction(tenant_id):
-            APITask.objects.create(
+            APITask.objects.update_or_create(
                 id=task_result_instance.task_id,
                 tenant_id=tenant_id,
-                task_runner_task=task_result_instance,
+                defaults={"task_runner_task": task_result_instance},
             )
         return result

api/src/backend/tasks/jobs/deletion.py

@@ -1,5 +1,5 @@
 from celery.utils.log import get_task_logger
-from django.db import transaction
+from django.db import DatabaseError
 
 from api.db_router import MainRouter
 from api.db_utils import batch_delete, rls_transaction
@@ -8,11 +8,12 @@ from api.models import Finding, Provider, Resource, Scan, ScanSummary, Tenant
 logger = get_task_logger(__name__)
 
 
-def delete_provider(pk: str):
+def delete_provider(tenant_id: str, pk: str):
     """
     Gracefully deletes an instance of a provider along with its related data.
 
     Args:
+        tenant_id (str): Tenant ID the resources belong to.
         pk (str): The primary key of the Provider instance to delete.
 
     Returns:
@@ -22,33 +23,31 @@ def delete_provider(pk: str):
     Raises:
         Provider.DoesNotExist: If no instance with the provided primary key exists.
     """
-    instance = Provider.all_objects.get(pk=pk)
-    deletion_summary = {}
+    with rls_transaction(tenant_id):
+        instance = Provider.all_objects.get(pk=pk)
+        deletion_summary = {}
+        deletion_steps = [
+            ("Scan Summaries", ScanSummary.all_objects.filter(scan__provider=instance)),
+            ("Findings", Finding.all_objects.filter(scan__provider=instance)),
+            ("Resources", Resource.all_objects.filter(provider=instance)),
+            ("Scans", Scan.all_objects.filter(provider=instance)),
+        ]
 
-    with transaction.atomic():
-        # Delete Scan Summaries
-        scan_summaries_qs = ScanSummary.all_objects.filter(scan__provider=instance)
-        _, scans_summ_summary = batch_delete(scan_summaries_qs)
-        deletion_summary.update(scans_summ_summary)
+    for step_name, queryset in deletion_steps:
+        try:
+            _, step_summary = batch_delete(tenant_id, queryset)
+            deletion_summary.update(step_summary)
+        except DatabaseError as db_error:
+            logger.error(f"Error deleting {step_name}: {db_error}")
+            raise
 
-        # Delete Findings
-        findings_qs = Finding.all_objects.filter(scan__provider=instance)
-        _, findings_summary = batch_delete(findings_qs)
-        deletion_summary.update(findings_summary)
-
-        # Delete Resources
-        resources_qs = Resource.all_objects.filter(provider=instance)
-        _, resources_summary = batch_delete(resources_qs)
-        deletion_summary.update(resources_summary)
-
-        # Delete Scans
-        scans_qs = Scan.all_objects.filter(provider=instance)
-        _, scans_summary = batch_delete(scans_qs)
-        deletion_summary.update(scans_summary)
-
-        provider_deleted_count, provider_summary = instance.delete()
+    try:
+        with rls_transaction(tenant_id):
+            _, provider_summary = instance.delete()
         deletion_summary.update(provider_summary)
-
+    except DatabaseError as db_error:
+        logger.error(f"Error deleting Provider: {db_error}")
+        raise
     return deletion_summary
 
 
@@ -66,9 +65,8 @@ def delete_tenant(pk: str):
     deletion_summary = {}
 
     for provider in Provider.objects.using(MainRouter.admin_db).filter(tenant_id=pk):
-        with rls_transaction(pk):
-            summary = delete_provider(provider.id)
-            deletion_summary.update(summary)
+        summary = delete_provider(pk, provider.id)
+        deletion_summary.update(summary)
 
     Tenant.objects.using(MainRouter.admin_db).filter(id=pk).delete()
 

api/src/backend/tasks/tasks.py

@@ -43,9 +43,10 @@ def check_provider_connection_task(provider_id: str):
     return check_provider_connection(provider_id=provider_id)
 
 
-@shared_task(base=RLSTask, name="provider-deletion", queue="deletion")
-@set_tenant
-def delete_provider_task(provider_id: str):
+@shared_task(
+    base=RLSTask, name="provider-deletion", queue="deletion", autoretry_for=(Exception,)
+)
+def delete_provider_task(provider_id: str, tenant_id: str):
     """
     Task to delete a specific Provider instance.
 
@@ -53,6 +54,7 @@ def delete_provider_task(provider_id: str):
 
     Args:
         provider_id (str): The primary key of the `Provider` instance to be deleted.
+        tenant_id (str): Tenant ID the provider belongs to.
 
     Returns:
         tuple: A tuple containing:
@@ -60,7 +62,7 @@ def delete_provider_task(provider_id: str):
             - A dictionary with the count of deleted instances per model,
               including related models if cascading deletes were triggered.
     """
-    return delete_provider(pk=provider_id)
+    return delete_provider(tenant_id=tenant_id, pk=provider_id)
 
 
 @shared_task(base=RLSTask, name="scan-perform", queue="scans")
@@ -174,7 +176,7 @@ def perform_scan_summary_task(tenant_id: str, scan_id: str):
     return aggregate_findings(tenant_id=tenant_id, scan_id=scan_id)
 
 
-@shared_task(name="tenant-deletion", queue="deletion")
+@shared_task(name="tenant-deletion", queue="deletion", autoretry_for=(Exception,))
 def delete_tenant_task(tenant_id: str):
     return delete_tenant(pk=tenant_id)
 

api/src/backend/tasks/tests/test_deletion.py

@@ -9,17 +9,19 @@ from api.models import Provider, Tenant
 class TestDeleteProvider:
     def test_delete_provider_success(self, providers_fixture):
         instance = providers_fixture[0]
-        result = delete_provider(instance.id)
+        tenant_id = str(instance.tenant_id)
+        result = delete_provider(tenant_id, instance.id)
 
         assert result
         with pytest.raises(ObjectDoesNotExist):
             Provider.objects.get(pk=instance.id)
 
-    def test_delete_provider_does_not_exist(self):
+    def test_delete_provider_does_not_exist(self, tenants_fixture):
+        tenant_id = str(tenants_fixture[0].id)
         non_existent_pk = "babf6796-cfcc-4fd3-9dcf-88d012247645"
 
         with pytest.raises(ObjectDoesNotExist):
-            delete_provider(non_existent_pk)
+            delete_provider(tenant_id, non_existent_pk)
 
 
 @pytest.mark.django_db