From fb914a2c908d1dc0b4aac36efb4db8d34fc4fbb6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adri=C3=A1n=20Jes=C3=BAs=20Pe=C3=B1a=20Rodr=C3=ADguez?= Date: Wed, 4 Jun 2025 15:01:47 +0200 Subject: [PATCH] revert: remove get_with_retry (#7932) --- api/CHANGELOG.md | 1 + api/src/backend/api/models.py | 40 ------------------------ api/src/backend/api/tests/test_models.py | 37 +--------------------- api/src/backend/api/v1/views.py | 8 ++--- 4 files changed, 6 insertions(+), 80 deletions(-) diff --git a/api/CHANGELOG.md b/api/CHANGELOG.md index dd3446550a..c840379feb 100644 --- a/api/CHANGELOG.md +++ b/api/CHANGELOG.md @@ -21,6 +21,7 @@ All notable changes to the **Prowler API** are documented in this file. ### Fixed - Fixed transaction persistence with RLS operations [(#7916)](https://github.com/prowler-cloud/prowler/pull/7916). +- Reverted the change `get_with_retry` to use the original `get` method for retrieving tasks [(#7932)](https://github.com/prowler-cloud/prowler/pull/7932). --- diff --git a/api/src/backend/api/models.py b/api/src/backend/api/models.py index b4a5c3a894..52e32b8c38 100644 --- a/api/src/backend/api/models.py +++ b/api/src/backend/api/models.py @@ -1,9 +1,7 @@ import json import re -import time from uuid import UUID, uuid4 -from config.env import env from cryptography.fernet import Fernet from django.conf import settings from django.contrib.auth.models import AbstractBaseUser @@ -354,42 +352,6 @@ class ProviderGroupMembership(RowLevelSecurityProtectedModel): resource_name = "provider_groups-provider" -class TaskManager(models.Manager): - def get_with_retry( - self, - id: str, - max_retries: int = None, - delay_seconds: float = None, - ): - """ - Retry fetching a Task by ID in case it hasn't been created yet. - - Args: - id (str): The Celery task ID (expected to match Task model PK). - max_retries (int, optional): Number of retry attempts. Defaults to env TASK_RETRY_ATTEMPTS or 5. - delay_seconds (float, optional): Delay between retries in seconds. Defaults to env TASK_RETRY_DELAY_SECONDS or 0.1. - - Returns: - Task: The retrieved Task instance. - - Raises: - Task.DoesNotExist: If the task is not found after all retries. - """ - max_retries = max_retries or env.int("TASK_RETRY_ATTEMPTS", default=5) - delay_seconds = delay_seconds or env.float( - "TASK_RETRY_DELAY_SECONDS", default=0.1 - ) - - for _attempt in range(max_retries): - try: - return self.get(id=id) - except self.model.DoesNotExist: - time.sleep(delay_seconds) - raise self.model.DoesNotExist( - f"Task with ID {id} not found after {max_retries} retries." - ) - - class Task(RowLevelSecurityProtectedModel): id = models.UUIDField(primary_key=True, default=uuid4, editable=False) inserted_at = models.DateTimeField(auto_now_add=True, editable=False) @@ -402,8 +364,6 @@ class Task(RowLevelSecurityProtectedModel): blank=True, ) - objects = TaskManager() - class Meta(RowLevelSecurityProtectedModel.Meta): db_table = "tasks" diff --git a/api/src/backend/api/tests/test_models.py b/api/src/backend/api/tests/test_models.py index de2ec7c59e..c2beeb3583 100644 --- a/api/src/backend/api/tests/test_models.py +++ b/api/src/backend/api/tests/test_models.py @@ -1,9 +1,6 @@ -import uuid -from unittest import mock - import pytest -from api.models import Resource, ResourceTag, Task +from api.models import Resource, ResourceTag @pytest.mark.django_db @@ -123,35 +120,3 @@ class TestResourceModel: # compliance={}, # ) # assert Finding.objects.filter(uid=long_uid).exists() - - -@pytest.mark.django_db -class TestTaskManager: - def test_get_with_retry_success(self): - task_id = uuid.uuid4() - call_counter = {"count": 0} - - def side_effect(*args, **kwargs): - if call_counter["count"] < 2: - call_counter["count"] += 1 - raise Task.DoesNotExist() - return Task(id=task_id) - - with mock.patch.object(Task.objects, "get", side_effect=side_effect): - task = Task.objects.get_with_retry( - task_id, max_retries=5, delay_seconds=0.01 - ) - - assert task.id == task_id - assert call_counter["count"] == 2 - - def test_get_with_retry_fail(self): - non_existent_id = uuid.uuid4() - - with mock.patch.object(Task.objects, "get", side_effect=Task.DoesNotExist): - with pytest.raises(Task.DoesNotExist) as excinfo: - Task.objects.get_with_retry( - non_existent_id, max_retries=3, delay_seconds=0.01 - ) - - assert str(non_existent_id) in str(excinfo.value) diff --git a/api/src/backend/api/v1/views.py b/api/src/backend/api/v1/views.py index 394913cb6c..99b551fec5 100644 --- a/api/src/backend/api/v1/views.py +++ b/api/src/backend/api/v1/views.py @@ -1092,7 +1092,7 @@ class ProviderViewSet(BaseRLSViewSet): task = check_provider_connection_task.delay( provider_id=pk, tenant_id=self.request.tenant_id ) - prowler_task = Task.objects.get_with_retry(id=task.id) + prowler_task = Task.objects.get(id=task.id) serializer = TaskSerializer(prowler_task) return Response( data=serializer.data, @@ -1115,7 +1115,7 @@ class ProviderViewSet(BaseRLSViewSet): task = delete_provider_task.delay( provider_id=pk, tenant_id=self.request.tenant_id ) - prowler_task = Task.objects.get_with_retry(id=task.id) + prowler_task = Task.objects.get(id=task.id) serializer = TaskSerializer(prowler_task) return Response( data=serializer.data, @@ -1495,7 +1495,7 @@ class ScanViewSet(BaseRLSViewSet): }, ) - prowler_task = Task.objects.get_with_retry(id=task.id) + prowler_task = Task.objects.get(id=task.id) scan.task_id = task.id scan.save(update_fields=["task_id"]) @@ -3123,7 +3123,7 @@ class ScheduleViewSet(BaseRLSViewSet): with transaction.atomic(): task = schedule_provider_scan(provider_instance) - prowler_task = Task.objects.get_with_retry(id=task.id) + prowler_task = Task.objects.get(id=task.id) self.response_serializer_class = TaskSerializer output_serializer = self.get_serializer(prowler_task)