feat(findings): Improve performance on /findings/metadata, /overviews and filters (#7690)

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

@@ -0,0 +1,61 @@
+from api.db_utils import rls_transaction
+from api.models import (
+    Resource,
+    ResourceFindingMapping,
+    ResourceScanSummary,
+    Scan,
+    StateChoices,
+)
+
+
+def backfill_resource_scan_summaries(tenant_id: str, scan_id: str):
+    with rls_transaction(tenant_id):
+        if ResourceScanSummary.objects.filter(
+            tenant_id=tenant_id, scan_id=scan_id
+        ).exists():
+            return {"status": "already backfilled"}
+
+    with rls_transaction(tenant_id):
+        if not Scan.objects.filter(
+            tenant_id=tenant_id,
+            id=scan_id,
+            state__in=(StateChoices.COMPLETED, StateChoices.FAILED),
+        ).exists():
+            return {"status": "scan is not completed"}
+
+        resource_ids_qs = (
+            ResourceFindingMapping.objects.filter(
+                tenant_id=tenant_id, finding__scan_id=scan_id
+            )
+            .values_list("resource_id", flat=True)
+            .distinct()
+        )
+
+        resource_ids = list(resource_ids_qs)
+
+        if not resource_ids:
+            return {"status": "no resources to backfill"}
+
+        resources_qs = Resource.objects.filter(
+            tenant_id=tenant_id, id__in=resource_ids
+        ).only("id", "service", "region", "type")
+
+        summaries = []
+        for resource in resources_qs.iterator():
+            summaries.append(
+                ResourceScanSummary(
+                    tenant_id=tenant_id,
+                    scan_id=scan_id,
+                    resource_id=str(resource.id),
+                    service=resource.service,
+                    region=resource.region,
+                    resource_type=resource.type,
+                )
+            )
+
+        for i in range(0, len(summaries), 500):
+            ResourceScanSummary.objects.bulk_create(
+                summaries[i : i + 500], ignore_conflicts=True
+            )
+
+    return {"status": "backfilled", "inserted": len(summaries)}

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

@@ -19,6 +19,7 @@ from api.models import (
     Finding,
     Provider,
     Resource,
+    ResourceScanSummary,
     ResourceTag,
     Scan,
     ScanSummary,
@@ -121,6 +122,7 @@ def perform_prowler_scan(
     check_status_by_region = {}
     exception = None
     unique_resources = set()
+    scan_resource_cache: set[tuple[str, str, str, str]] = set()
     start_time = time.time()
 
     with rls_transaction(tenant_id):
@@ -295,6 +297,16 @@ def perform_prowler_scan(
                     continue
                 region_dict[finding.check_id] = finding.status.value
 
+                # Update scan resource summaries
+                scan_resource_cache.add(
+                    (
+                        str(resource_instance.id),
+                        resource_instance.service,
+                        resource_instance.region,
+                        resource_instance.type,
+                    )
+                )
+
             # Update scan progress
             with rls_transaction(tenant_id):
                 scan_instance.progress = progress
@@ -314,66 +326,90 @@ def perform_prowler_scan(
             scan_instance.unique_resource_count = len(unique_resources)
             scan_instance.save()
 
-    if exception is None:
-        try:
-            regions = prowler_provider.get_regions()
-        except AttributeError:
-            regions = set()
-
-        compliance_template = PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE[
-            provider_instance.provider
-        ]
-        compliance_overview_by_region = {
-            region: deepcopy(compliance_template) for region in regions
-        }
-
-        for region, check_status in check_status_by_region.items():
-            compliance_data = compliance_overview_by_region.setdefault(
-                region, deepcopy(compliance_template)
-            )
-            for check_name, status in check_status.items():
-                generate_scan_compliance(
-                    compliance_data,
-                    provider_instance.provider,
-                    check_name,
-                    status,
-                )
-
-        # Prepare compliance overview objects
-        compliance_overview_objects = []
-        for region, compliance_data in compliance_overview_by_region.items():
-            for compliance_id, compliance in compliance_data.items():
-                compliance_overview_objects.append(
-                    ComplianceOverview(
-                        tenant_id=tenant_id,
-                        scan=scan_instance,
-                        region=region,
-                        compliance_id=compliance_id,
-                        framework=compliance["framework"],
-                        version=compliance["version"],
-                        description=compliance["description"],
-                        requirements=compliance["requirements"],
-                        requirements_passed=compliance["requirements_status"]["passed"],
-                        requirements_failed=compliance["requirements_status"]["failed"],
-                        requirements_manual=compliance["requirements_status"]["manual"],
-                        total_requirements=compliance["total_requirements"],
-                    )
-                )
-        try:
-            with rls_transaction(tenant_id):
-                ComplianceOverview.objects.bulk_create(
-                    compliance_overview_objects, batch_size=100
-                )
-        except Exception as overview_exception:
-            import sentry_sdk
-
-            sentry_sdk.capture_exception(overview_exception)
-            logger.error(
-                f"Error storing compliance overview for scan {scan_id}: {overview_exception}"
-            )
     if exception is not None:
         raise exception
 
+    try:
+        regions = prowler_provider.get_regions()
+    except AttributeError:
+        regions = set()
+
+    compliance_template = PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE[
+        provider_instance.provider
+    ]
+    compliance_overview_by_region = {
+        region: deepcopy(compliance_template) for region in regions
+    }
+
+    for region, check_status in check_status_by_region.items():
+        compliance_data = compliance_overview_by_region.setdefault(
+            region, deepcopy(compliance_template)
+        )
+        for check_name, status in check_status.items():
+            generate_scan_compliance(
+                compliance_data,
+                provider_instance.provider,
+                check_name,
+                status,
+            )
+
+    # Prepare compliance overview objects
+    compliance_overview_objects = []
+    for region, compliance_data in compliance_overview_by_region.items():
+        for compliance_id, compliance in compliance_data.items():
+            compliance_overview_objects.append(
+                ComplianceOverview(
+                    tenant_id=tenant_id,
+                    scan=scan_instance,
+                    region=region,
+                    compliance_id=compliance_id,
+                    framework=compliance["framework"],
+                    version=compliance["version"],
+                    description=compliance["description"],
+                    requirements=compliance["requirements"],
+                    requirements_passed=compliance["requirements_status"]["passed"],
+                    requirements_failed=compliance["requirements_status"]["failed"],
+                    requirements_manual=compliance["requirements_status"]["manual"],
+                    total_requirements=compliance["total_requirements"],
+                )
+            )
+    try:
+        with rls_transaction(tenant_id):
+            ComplianceOverview.objects.bulk_create(
+                compliance_overview_objects, batch_size=500
+            )
+    except Exception as overview_exception:
+        import sentry_sdk
+
+        sentry_sdk.capture_exception(overview_exception)
+        logger.error(
+            f"Error storing compliance overview for scan {scan_id}: {overview_exception}"
+        )
+
+    try:
+        resource_scan_summaries = [
+            ResourceScanSummary(
+                tenant_id=tenant_id,
+                scan_id=scan_id,
+                resource_id=resource_id,
+                service=service,
+                region=region,
+                resource_type=resource_type,
+            )
+            for resource_id, service, region, resource_type in scan_resource_cache
+        ]
+        with rls_transaction(tenant_id):
+            ResourceScanSummary.objects.bulk_create(
+                resource_scan_summaries, batch_size=500, ignore_conflicts=True
+            )
+    except Exception as filter_exception:
+        import sentry_sdk
+
+        sentry_sdk.capture_exception(filter_exception)
+        logger.error(
+            f"Error storing filter values for scan {scan_id}: {filter_exception}"
+        )
+
     serializer = ScanTaskSerializer(instance=scan_instance)
     return serializer.data
 

api/src/backend/tasks/tasks.py

@@ -7,6 +7,7 @@ from celery.utils.log import get_task_logger
 from config.celery import RLSTask
 from config.django.base import DJANGO_FINDINGS_BATCH_SIZE, DJANGO_TMP_OUTPUT_DIRECTORY
 from django_celery_beat.models import PeriodicTask
+from tasks.jobs.backfill import backfill_resource_scan_summaries
 from tasks.jobs.connection import check_provider_connection
 from tasks.jobs.deletion import delete_provider, delete_tenant
 from tasks.jobs.export import (
@@ -358,3 +359,15 @@ def generate_outputs(scan_id: str, provider_id: str, tenant_id: str):
     Scan.all_objects.filter(id=scan_id).update(output_location=final_location)
     logger.info(f"Scan outputs at {final_location}")
     return {"upload": did_upload}
+
+
+@shared_task(name="backfill-scan-resource-summaries", queue="backfill")
+def backfill_scan_resource_summaries_task(tenant_id: str, scan_id: str):
+    """
+    Tries to backfill the resource scan summaries table for a given scan.
+
+    Args:
+        tenant_id (str): The tenant identifier.
+        scan_id (str): The scan identifier.
+    """
+    return backfill_resource_scan_summaries(tenant_id=tenant_id, scan_id=scan_id)

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

@@ -0,0 +1,79 @@
+from uuid import uuid4
+
+import pytest
+from tasks.jobs.backfill import backfill_resource_scan_summaries
+
+from api.models import ResourceScanSummary, Scan, StateChoices
+
+
+@pytest.mark.django_db
+class TestBackfillResourceScanSummaries:
+    @pytest.fixture(scope="function")
+    def resource_scan_summary_data(self, scans_fixture):
+        scan = scans_fixture[0]
+        return ResourceScanSummary.objects.create(
+            tenant_id=scan.tenant_id,
+            scan_id=scan.id,
+            resource_id=str(uuid4()),
+            service="aws",
+            region="us-east-1",
+            resource_type="instance",
+        )
+
+    @pytest.fixture(scope="function")
+    def get_not_completed_scans(self, providers_fixture):
+        provider_id = providers_fixture[0].id
+        tenant_id = providers_fixture[0].tenant_id
+        scan_1 = Scan.objects.create(
+            tenant_id=tenant_id,
+            trigger=Scan.TriggerChoices.MANUAL,
+            state=StateChoices.EXECUTING,
+            provider_id=provider_id,
+        )
+        scan_2 = Scan.objects.create(
+            tenant_id=tenant_id,
+            trigger=Scan.TriggerChoices.MANUAL,
+            state=StateChoices.AVAILABLE,
+            provider_id=provider_id,
+        )
+        return scan_1, scan_2
+
+    def test_already_backfilled(self, resource_scan_summary_data):
+        tenant_id = resource_scan_summary_data.tenant_id
+        scan_id = resource_scan_summary_data.scan_id
+
+        result = backfill_resource_scan_summaries(tenant_id, scan_id)
+
+        assert result == {"status": "already backfilled"}
+
+    def test_not_completed_scan(self, get_not_completed_scans):
+        for scan_instance in get_not_completed_scans:
+            tenant_id = scan_instance.tenant_id
+            scan_id = scan_instance.id
+            result = backfill_resource_scan_summaries(tenant_id, scan_id)
+
+            assert result == {"status": "scan is not completed"}
+
+    def test_successful_backfill_inserts_one_summary(
+        self, resources_fixture, findings_fixture
+    ):
+        tenant_id = findings_fixture[0].tenant_id
+        scan_id = findings_fixture[0].scan_id
+
+        # This scan affects the first two resources
+        resources = resources_fixture[:2]
+
+        result = backfill_resource_scan_summaries(tenant_id, scan_id)
+        assert result == {"status": "backfilled", "inserted": len(resources)}
+
+        # Verify correct values
+        summaries = ResourceScanSummary.objects.filter(
+            tenant_id=tenant_id, scan_id=scan_id
+        )
+        assert summaries.count() == len(resources)
+        for resource in resources:
+            summary = summaries.get(resource_id=resource.id)
+            assert summary.resource_id == resource.id
+            assert summary.service == resource.service
+            assert summary.region == resource.region
+            assert summary.resource_type == resource.type