diff --git a/src/backend/api/migrations/0001_initial.py b/src/backend/api/migrations/0001_initial.py index 1b6a4f8d92..6a6e389d3d 100644 --- a/src/backend/api/migrations/0001_initial.py +++ b/src/backend/api/migrations/0001_initial.py @@ -864,7 +864,7 @@ class Migration(migrations.Migration): migrations.AddConstraint( model_name="resource", constraint=models.UniqueConstraint( - fields=("tenant_id", "provider_id", "uid", "region"), + fields=("tenant_id", "provider_id", "uid"), name="unique_resources_by_provider", ), ), diff --git a/src/backend/api/models.py b/src/backend/api/models.py index 2196073b23..e71b2b57cc 100644 --- a/src/backend/api/models.py +++ b/src/backend/api/models.py @@ -432,7 +432,7 @@ class Resource(RowLevelSecurityProtectedModel): constraints = [ models.UniqueConstraint( - fields=("tenant_id", "provider_id", "uid", "region"), + fields=("tenant_id", "provider_id", "uid"), name="unique_resources_by_provider", ), RowLevelSecurityConstraint( diff --git a/src/backend/tasks/jobs/scan.py b/src/backend/tasks/jobs/scan.py index 1380171d2a..e023cc7b0b 100644 --- a/src/backend/tasks/jobs/scan.py +++ b/src/backend/tasks/jobs/scan.py @@ -108,14 +108,22 @@ def _store_resources( """ with tenant_transaction(tenant_id): - resource_instance, _ = Resource.objects.get_or_create( + resource_instance, created = Resource.objects.get_or_create( tenant_id=tenant_id, provider=provider_instance, uid=finding.resource_uid, - region=finding.region, - service=finding.service_name, - type=finding.resource_type, + defaults={ + "region": finding.region, + "service": finding.service_name, + "type": finding.resource_type, + }, ) + + if not created: + resource_instance.region = finding.region + resource_instance.service = finding.service_name + resource_instance.type = finding.resource_type + resource_instance.save() with tenant_transaction(tenant_id): tags = [ ResourceTag.objects.get_or_create( diff --git a/src/backend/tasks/tests/test_scan.py b/src/backend/tasks/tests/test_scan.py index d98c037375..23509233f8 100644 --- a/src/backend/tasks/tests/test_scan.py +++ b/src/backend/tasks/tests/test_scan.py @@ -3,7 +3,11 @@ from unittest.mock import patch, MagicMock import pytest from api.models import StateChoices, StatusChoices, Severity, Finding, Resource -from tasks.jobs.scan import perform_prowler_scan, _create_finding_delta +from tasks.jobs.scan import ( + perform_prowler_scan, + _create_finding_delta, + _store_resources, +) @pytest.mark.django_db @@ -112,3 +116,142 @@ class TestPerformScan: ) def test_create_finding_delta(self, last_status, new_status, expected_delta): assert _create_finding_delta(last_status, new_status) == expected_delta + + @patch("tasks.jobs.scan.ResourceTag.objects.get_or_create") + @patch("tasks.jobs.scan.Resource.objects.get_or_create") + @patch("api.db_utils.tenant_transaction") + def test_store_resources_new_resource( + self, + mock_tenant_transaction, + mock_get_or_create_resource, + mock_get_or_create_tag, + ): + tenant_id = "tenant123" + provider_instance = MagicMock() + provider_instance.id = "provider456" + + finding = MagicMock() + finding.resource_uid = "resource_uid_123" + finding.region = "us-west-1" + finding.service_name = "service_name" + finding.resource_type = "resource_type" + finding.resource_tags = {"tag1": "value1", "tag2": "value2"} + + resource_instance = MagicMock() + mock_get_or_create_resource.return_value = (resource_instance, True) + tag_instance = MagicMock() + mock_get_or_create_tag.return_value = (tag_instance, True) + + resource, resource_uid_tuple = _store_resources( + finding, tenant_id, provider_instance + ) + + mock_get_or_create_resource.assert_called_once_with( + tenant_id=tenant_id, + provider=provider_instance, + uid=finding.resource_uid, + defaults={ + "region": finding.region, + "service": finding.service_name, + "type": finding.resource_type, + }, + ) + + assert resource == resource_instance + assert resource_uid_tuple == (resource_instance.uid, resource_instance.region) + + @patch("tasks.jobs.scan.ResourceTag.objects.get_or_create") + @patch("tasks.jobs.scan.Resource.objects.get_or_create") + @patch("api.db_utils.tenant_transaction") + def test_store_resources_existing_resource( + self, + mock_tenant_transaction, + mock_get_or_create_resource, + mock_get_or_create_tag, + ): + tenant_id = "tenant123" + provider_instance = MagicMock() + provider_instance.id = "provider456" + + finding = MagicMock() + finding.resource_uid = "resource_uid_123" + finding.region = "us-west-2" + finding.service_name = "new_service" + finding.resource_type = "new_type" + finding.resource_tags = {"tag1": "value1", "tag2": "value2"} + + resource_instance = MagicMock() + resource_instance.region = "us-west-1" + resource_instance.service = "old_service" + resource_instance.type = "old_type" + mock_get_or_create_resource.return_value = (resource_instance, False) + + tag_instance = MagicMock() + mock_get_or_create_tag.return_value = (tag_instance, True) + + resource, resource_uid_tuple = _store_resources( + finding, tenant_id, provider_instance + ) + + mock_get_or_create_resource.assert_called_once_with( + tenant_id=tenant_id, + provider=provider_instance, + uid=finding.resource_uid, + defaults={ + "region": finding.region, + "service": finding.service_name, + "type": finding.resource_type, + }, + ) + + assert resource_instance.region == finding.region + assert resource_instance.service == finding.service_name + assert resource_instance.type == finding.resource_type + assert resource == resource_instance + assert resource_uid_tuple == (resource_instance.uid, resource_instance.region) + + @patch("tasks.jobs.scan.ResourceTag.objects.get_or_create") + @patch("tasks.jobs.scan.Resource.objects.get_or_create") + @patch("api.db_utils.tenant_transaction") + def test_store_resources_with_tags( + self, + mock_tenant_transaction, + mock_get_or_create_resource, + mock_get_or_create_tag, + ): + tenant_id = "tenant123" + provider_instance = MagicMock() + provider_instance.id = "provider456" + + finding = MagicMock() + finding.resource_uid = "resource_uid_123" + finding.region = "us-west-1" + finding.service_name = "service_name" + finding.resource_type = "resource_type" + finding.resource_tags = {"tag1": "value1", "tag2": "value2"} + + resource_instance = MagicMock() + mock_get_or_create_resource.return_value = (resource_instance, True) + tag_instance_1 = MagicMock() + tag_instance_2 = MagicMock() + mock_get_or_create_tag.side_effect = [ + (tag_instance_1, True), + (tag_instance_2, True), + ] + + resource, resource_uid_tuple = _store_resources( + finding, tenant_id, provider_instance + ) + + mock_get_or_create_tag.assert_any_call( + tenant_id=tenant_id, key="tag1", value="value1" + ) + mock_get_or_create_tag.assert_any_call( + tenant_id=tenant_id, key="tag2", value="value2" + ) + resource_instance.upsert_or_delete_tags.assert_called_once_with( + tags=[tag_instance_1, tag_instance_2] + ) + + assert resource == resource_instance + assert resource_uid_tuple == (resource_instance.uid, resource_instance.region)