Prevent committing partial advisory data to db

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

vulnerabilities/pipes/advisory.py

@@ -134,6 +134,7 @@ def insert_advisory(advisory: AdvisoryData, pipeline_id: str, logger: Callable =
     return advisory_obj
 
 
+@transaction.atomic
 def insert_advisory_v2(
     advisory: AdvisoryData,
     pipeline_id: str,

vulnerabilities/tests/pipes/test_vulnerablecode_importer_pipeline_v2.py

@@ -10,11 +10,17 @@
 import logging
 from datetime import datetime
 from datetime import timedelta
+from unittest.mock import patch
 
 import pytest
+from packageurl import PackageURL
+from univers.version_range import VersionRange
 
 from vulnerabilities.importer import AdvisoryData
+from vulnerabilities.importer import AffectedPackageV2
 from vulnerabilities.models import AdvisoryV2
+from vulnerabilities.models import ImpactedPackage
+from vulnerabilities.models import PackageV2
 from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipelineV2
 
 
@@ -40,7 +46,18 @@ def dummy_advisory():
         references_v2=[],
         severities=[],
         weaknesses=[],
-        affected_packages=[],
+        affected_packages=[
+            AffectedPackageV2(
+                package=PackageURL.from_string("pkg:npm/foobar"),
+                affected_version_range=VersionRange.from_string("vers:npm/<=1.2.3"),
+                fixed_version_range=VersionRange.from_string("vers:npm/1.2.4"),
+            ),
+            AffectedPackageV2(
+                package=PackageURL.from_string("pkg:npm/foobar"),
+                affected_version_range=VersionRange.from_string("vers:npm/<=3.2.3"),
+                fixed_version_range=VersionRange.from_string("vers:npm/3.2.4"),
+            ),
+        ],
         advisory_id="ADV-123",
         date_published=datetime.now() - timedelta(days=10),
         url="https://example.com/advisory/1",
@@ -60,3 +77,19 @@ def test_collect_and_store_advisories(dummy_importer):
     assert len(dummy_importer.log_messages) >= 2
     assert "Successfully collected" in dummy_importer.log_messages[-1][1]
     assert AdvisoryV2.objects.count() == 1
+
+
+@pytest.mark.django_db
+@patch("vulnerabilities.pipes.advisory.get_exact_purls_v2", side_effect=Exception("error"))
+def test_advisory_import_atomicity_no_partial_adv_import(mock_exception, dummy_importer):
+    dummy_importer.collect_and_store_advisories()
+    assert AdvisoryV2.objects.count() == 0
+    assert ImpactedPackage.objects.count() == 0
+
+
+@pytest.mark.django_db
+def test_advisory_import_atomicity(dummy_importer):
+    dummy_importer.collect_and_store_advisories()
+    assert AdvisoryV2.objects.count() == 1
+    assert ImpactedPackage.objects.count() == 2
+    assert PackageV2.objects.count() == 4