From 7fa57ba3e2b89b529835c81e2352abb8deadfb5c Mon Sep 17 00:00:00 2001 From: Prowler Bot Date: Mon, 24 Mar 2025 14:58:39 +0100 Subject: [PATCH] ref(providers): Refactor provider deletion functions (#7351) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Víctor Fernández Poyatos --- api/CHANGELOG.md | 7 +++ api/src/backend/api/db_utils.py | 33 ++++++++---- api/src/backend/api/signals.py | 7 ++- api/src/backend/api/specs/v1.yaml | 2 +- api/src/backend/api/tests/test_db_utils.py | 5 +- api/src/backend/api/v1/views.py | 4 +- api/src/backend/config/celery.py | 4 +- api/src/backend/tasks/jobs/deletion.py | 56 ++++++++++---------- api/src/backend/tasks/tasks.py | 12 +++-- api/src/backend/tasks/tests/test_deletion.py | 8 +-- 10 files changed, 82 insertions(+), 56 deletions(-) diff --git a/api/CHANGELOG.md b/api/CHANGELOG.md index 9a2ee793a5..e82bac48f8 100644 --- a/api/CHANGELOG.md +++ b/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 diff --git a/api/src/backend/api/db_utils.py b/api/src/backend/api/db_utils.py index 1bcf14209c..02b4f9b455 100644 --- a/api/src/backend/api/db_utils.py +++ b/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 diff --git a/api/src/backend/api/signals.py b/api/src/backend/api/signals.py index 44c2e0b4fd..d9dc0a72d9 100644 --- a/api/src/backend/api/signals.py +++ b/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) diff --git a/api/src/backend/api/specs/v1.yaml b/api/src/backend/api/specs/v1.yaml index 805fff1666..a65c97f7e0 100644 --- a/api/src/backend/api/specs/v1.yaml +++ b/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. diff --git a/api/src/backend/api/tests/test_db_utils.py b/api/src/backend/api/tests/test_db_utils.py index 6c2364600a..e22b1417bc 100644 --- a/api/src/backend/api/tests/test_db_utils.py +++ b/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} diff --git a/api/src/backend/api/v1/views.py b/api/src/backend/api/v1/views.py index 1851bf5426..9fff8ce853 100644 --- a/api/src/backend/api/v1/views.py +++ b/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( diff --git a/api/src/backend/config/celery.py b/api/src/backend/config/celery.py index 4486a0bf97..4c667b7bc5 100644 --- a/api/src/backend/config/celery.py +++ b/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 diff --git a/api/src/backend/tasks/jobs/deletion.py b/api/src/backend/tasks/jobs/deletion.py index 5ca08e70bb..d72b8de40e 100644 --- a/api/src/backend/tasks/jobs/deletion.py +++ b/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() diff --git a/api/src/backend/tasks/tasks.py b/api/src/backend/tasks/tasks.py index 05c1919f17..c9b8d1ab2a 100644 --- a/api/src/backend/tasks/tasks.py +++ b/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) diff --git a/api/src/backend/tasks/tests/test_deletion.py b/api/src/backend/tasks/tests/test_deletion.py index 00d218e192..81cdb44daa 100644 --- a/api/src/backend/tasks/tests/test_deletion.py +++ b/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