Merge pull request #2289 from aboutcode-org/package-bulk-create-fix

fix: ignore conflicts while bulk creating v2 packages

Author: keshav-space

vulnerabilities/models.py

@@ -3399,19 +3399,12 @@ def get_or_create_from_purl(self, purl: Union[PackageURL, str]):
         return package, is_created
 
     def bulk_get_or_create_from_purls(self, purls: List[Union[PackageURL, str]]):
-        """
-        Return new or existing Packages given ``purls`` list of PackageURL object or PURL string.
-        """
-        purl_strings = [str(p) for p in purls]
-        existing_packages = PackageV2.objects.filter(package_url__in=purl_strings)
-        existing_purls = set(existing_packages.values_list("package_url", flat=True))
+        """Return queryset of Packages for a list of PURLs, bulk create any that do not already exist."""
 
-        all_packages = list(existing_packages)
         packages_to_create = []
-        for purl in purls:
-            if str(purl) in existing_purls:
-                continue
+        normalize_purls = []
 
+        for purl in purls:
             purl_dict = purl_to_dict(purl)
             purl = PackageURL(**purl_dict)
 
@@ -3422,16 +3415,16 @@ def bulk_get_or_create_from_purls(self, purls: List[Union[PackageURL, str]]):
             purl_dict["package_url"] = str(normalized)
             purl_dict["plain_package_url"] = str(utils.plain_purl(normalized))
 
+            normalize_purls.append(str(normalized))
             packages_to_create.append(PackageV2(**purl_dict))
 
         try:
-            new_packages = PackageV2.objects.bulk_create(packages_to_create)
+            PackageV2.objects.bulk_create(packages_to_create, ignore_conflicts=True)
         except Exception as e:
             logging.error(f"Error creating PackageV2: {e} \n {traceback_format_exc()}")
             return []
 
-        all_packages.extend(new_packages)
-        return all_packages
+        return PackageV2.objects.filter(package_url__in=normalize_purls)
 
     def only_vulnerable(self):
         return self._vulnerable(True)

vulnerabilities/pipelines/v2_improvers/unfurl_version_range.py

@@ -166,7 +166,10 @@ def bulk_create_with_m2m(purls, impact, relation, logger):
 
     affected_packages_v2 = PackageV2.objects.bulk_get_or_create_from_purls(purls=purls)
 
-    affected_packages_v2[-1].calculate_version_rank
+    if not affected_packages_v2.exists():
+        return 0
+
+    affected_packages_v2.first().calculate_version_rank
 
     relations = [
         relation(impacted_package=impact, package=package) for package in affected_packages_v2

vulnerabilities/tests/test_models.py

@@ -34,6 +34,7 @@
 from vulnerabilities.models import AdvisorySeverity
 from vulnerabilities.models import Alias
 from vulnerabilities.models import Package
+from vulnerabilities.models import PackageV2
 from vulnerabilities.models import Patch
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.severity_systems import CVSSV3
@@ -827,3 +828,38 @@ def test_advisoryv2_duplication_data(self):
         result = models.AdvisoryV2.objects.count()
 
         self.assertEqual(result, 2)
+
+
+class TestPackageV2BulkCreate(DjangoTestCase):
+    def setUp(self):
+        PackageV2.objects.get_or_create_from_purl(
+            "pkg:deb/ubuntu/linux@6.17.0-19.19?arch=source&distro=questing"
+        )
+        PackageV2.objects.get_or_create_from_purl("pkg:pypi/foo@1.2.3")
+        PackageV2.objects.get_or_create_from_purl("pkg:npm/foobar@3.2.3")
+        PackageV2.objects.get_or_create_from_purl("pkg:maven/foo@1.2.3")
+        PackageV2.objects.get_or_create_from_purl(
+            "pkg:deb/ubuntu/linux@6.17.0-4.4?arch=source&distro=questing"
+        )
+
+    def test_package_bulk_get_or_create_from_purls(self):
+        purls = [
+            "pkg:npm/foo@1.2.3",
+            "pkg:pypi/foo@1.2.3",
+            "pkg:npm/foobar@3.2.3",
+            "pkg:maven/foo@1.2.3",
+            "pkg:nuget/foo@1.2.3",
+            "pkg:deb/ubuntu/linux@6.17.0-22.22?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-20.20?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-14.14?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-12.12?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-8.8?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-7.7?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-6.6?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-5.5?arch=source&distro=questing",
+            "pkg:deb/ubuntu/linux@6.17.0-4.4?arch=source&distro=questing",
+        ]
+        result_qs = PackageV2.objects.bulk_get_or_create_from_purls(purls)
+        result = [p.package_url for p in result_qs]
+
+        self.assertCountEqual(result, purls)