Bulk load advisories for grouping

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

Author: TG1999

vulnerabilities/api_v3.py

@@ -10,7 +10,8 @@
 from collections import defaultdict
 from urllib.parse import urlencode
 
-from django.contrib.postgres.aggregates import ArrayAgg, JSONBAgg
+from django.contrib.postgres.aggregates import ArrayAgg
+from django.contrib.postgres.aggregates import JSONBAgg
 from django.db.models import Exists
 from django.db.models import F
 from django.db.models import Max
@@ -736,11 +737,9 @@ def get_impacts_bulk(packages):
     impact_map = defaultdict(dict)
 
     for impact in impacts:
-        impact_map[
-            impact["package_id"]
-        ][
-            impact["impacted_package__advisory__avid"]
-        ] = impact["fixed_by_packages"] or []
+        impact_map[impact["package_id"]][impact["impacted_package__advisory__avid"]] = (
+            impact["fixed_by_packages"] or []
+        )
 
     return impact_map
 

vulnerabilities/models.py

@@ -2963,6 +2963,20 @@ def latest_affecting_advisories_for_purls(self, purls):
         )
         return self.filter(id__in=Subquery(adv_ids))
 
+    def latest_affecting_advisory_purls_pairs(self, purls):
+        return (
+            ImpactedPackage.objects.filter(
+                affecting_packages__package_url__in=purls,
+                advisory__is_latest=True,
+                advisory___all_impacts_unfurled=True,
+            )
+            .values_list(
+                "affecting_packages__package_url",
+                "advisory_id",
+            )
+            .distinct()
+        )
+
     def latest_affecting_advisories_for_packages(self, packages):
         adv_ids = ImpactedPackage.objects.filter(
             affecting_packages__in=packages,
@@ -2997,6 +3011,20 @@ def latest_fixed_by_advisories_for_purls(self, purls):
 
         return self.filter(id__in=Subquery(adv_ids))
 
+    def latest_fixed_by_advisory_purls_pairs(self, purls):
+        return (
+            ImpactedPackage.objects.filter(
+                fixed_by_packages__package_url__in=purls,
+                advisory__is_latest=True,
+                advisory___all_impacts_unfurled=True,
+            )
+            .values_list(
+                "fixed_by_packages__package_url",
+                "advisory_id",
+            )
+            .distinct()
+        )
+
     def latest_advisories_for_purls(self, purls):
         adv_ids = (
             ImpactedPackage.objects.filter(

vulnerabilities/pipelines/v2_improvers/mark_unfurl_version_range.py

@@ -7,6 +7,8 @@
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
 
+from collections import defaultdict
+
 from django.db import transaction
 from django.db.models import Exists
 from django.db.models import Min
@@ -21,6 +23,7 @@
 from vulnerabilities.models import PipelineSchedule
 from vulnerabilities.pipelines import VulnerableCodePipeline
 from vulnerabilities.pipes.group_advisories import group_advisory_for_package
+from vulnerabilities.pipes.group_advisories import group_single_package_with_provided_advisories
 from vulnerabilities.pipes.risk_score import compute_package_risk_score_bulk
 from vulnerabilities.utils import TYPES_WITH_MULTIPLE_IMPORTERS
 
@@ -151,7 +154,46 @@ def complete_advisories_import(advisory_ids, successful_advisory_ids=[]):
 
     group_package_ids = affecting_package_ids | fixed_by_package_ids
 
-    for package in PackageV2.objects.filter(
+    packages = PackageV2.objects.filter(
         id__in=group_package_ids, type__in=TYPES_WITH_MULTIPLE_IMPORTERS
-    ).iterator(chunk_size=2000):
-        group_advisory_for_package(package)
+    ).only("package_url", "id")
+
+    group_advisories_for_packages_bulk_marking(packages)
+
+
+def group_advisories_for_packages_bulk_marking(packages):
+    purls = [package.package_url for package in packages]
+
+    affecting_pairs = AdvisoryV2.objects.latest_affecting_advisory_purls_pairs(purls)
+
+    fixed_pairs = AdvisoryV2.objects.latest_fixed_by_advisory_purls_pairs(purls)
+
+    affecting_ids = {adv_id for _, adv_id in affecting_pairs}
+    fixed_ids = {adv_id for _, adv_id in fixed_pairs}
+
+    all_adv_ids = affecting_ids | fixed_ids
+
+    advisories = AdvisoryV2.objects.filter(id__in=all_adv_ids).prefetch_related(
+        "aliases",
+        "impacted_packages__affecting_packages",
+        "impacted_packages__fixed_by_packages",
+    )
+
+    advisory_map = {a.id: a for a in advisories}
+
+    affecting_by_purl = defaultdict(list)
+
+    for purl, advisory_id in affecting_pairs:
+        affecting_by_purl[purl].append(advisory_map[advisory_id])
+
+    fixed_by_purl = defaultdict(list)
+
+    for purl, advisory_id in fixed_pairs:
+        fixed_by_purl[purl].append(advisory_map[advisory_id])
+
+    for package in packages:
+        group_single_package_with_provided_advisories(
+            package=package,
+            affecting_advisories=affecting_by_purl.get(package.purl, []),
+            fixed_by_advisories=fixed_by_purl.get(package.purl, []),
+        )

vulnerabilities/pipes/group_advisories.py

@@ -130,17 +130,25 @@ def group_advisory_for_package(package, logger=None):
     )
 
     try:
-        affected_groups: List[Group] = merge_advisories(affecting_advisories, package)
-        fixed_by_groups: List[Group] = merge_advisories(fixed_by_advisories, package)
-        delete_and_save_advisory_set(affected_groups, package, relation="affecting")
-        delete_and_save_advisory_set(fixed_by_groups, package, relation="fixing")
+        group_single_package_with_provided_advisories(
+            package, affecting_advisories, fixed_by_advisories
+        )
         logger(f"Successfully rebuilt advisory sets for package {package.purl}")
     except Exception as e:
         if logger:
             logger(f"Failed rebuilding advisory sets for package {package.purl}: {e!r}")
         return
 
 
+def group_single_package_with_provided_advisories(
+    package, affecting_advisories, fixed_by_advisories
+):
+    affected_groups: List[Group] = merge_advisories(affecting_advisories, package)
+    fixed_by_groups: List[Group] = merge_advisories(fixed_by_advisories, package)
+    delete_and_save_advisory_set(affected_groups, package, relation="affecting")
+    delete_and_save_advisory_set(fixed_by_groups, package, relation="fixing")
+
+
 def compute_advisory_content_hash(adv, version_less_purl_str: str):
     """
     Compute a content hash for an advisory.