Add severities in the prefetch and optimize the prefetching process for compute_and_store_package_risk_score

Signed-off-by: ziad hany <[email protected]>

Author: ziadhany

vulnerabilities/pipelines/compute_package_risk.py

@@ -8,6 +8,7 @@
 #
 
 from aboutcode.pipeline import LoopProgress
+from django.db.models import Prefetch
 
 from vulnerabilities.models import Package
 from vulnerabilities.models import Vulnerability
@@ -38,6 +39,7 @@ def compute_and_store_vulnerability_risk_score(self):
             affectedbypackagerelatedvulnerability__isnull=False
         ).prefetch_related(
             "references",
+            "severities",
             "exploits",
         )
 
@@ -77,8 +79,15 @@ def compute_and_store_vulnerability_risk_score(self):
 
     def compute_and_store_package_risk_score(self):
         affected_packages = (
-            Package.objects.filter(affected_by_vulnerabilities__isnull=False).only("id").distinct()
-        )
+            Package.objects.filter(affected_by_vulnerabilities__isnull=False)
+            .only("id")
+            .prefetch_related(
+                Prefetch(
+                    "affectedbypackagerelatedvulnerability_set__vulnerability",
+                    queryset=Vulnerability.objects.only("weighted_severity", "exploitability"),
+                ),
+            )
+        ).distinct()
 
         self.log(f"Calculating risk for {affected_packages.count():,d} affected package records")
 

vulnerabilities/risk.py

@@ -8,10 +8,6 @@
 #
 from urllib.parse import urlparse
 
-from django.db.models import Prefetch
-
-from vulnerabilities.models import AffectedByPackageRelatedVulnerability
-from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference
 from vulnerabilities.severity_systems import EPSS
 from vulnerabilities.weight_config import WEIGHT_CONFIG
@@ -107,18 +103,10 @@ def compute_package_risk(package):
     Calculate the risk for a package by iterating over all vulnerabilities that affects this package
     and determining the associated risk.
     """
-
     result = []
-    affected_pkg_related_vul = AffectedByPackageRelatedVulnerability.objects.filter(
-        package=package
-    ).prefetch_related(
-        Prefetch(
-            "vulnerability",
-            queryset=Vulnerability.objects.only("weighted_severity", "exploitability"),
-        )
-    )
-    for pkg_related_vul in affected_pkg_related_vul:
-        if risk := pkg_related_vul.vulnerability.risk_score:
+    vulnerabilities = package.vulnerabilities.all()
+    for vulnerability in vulnerabilities:
+        if risk := vulnerability.risk_score:
             result.append(float(risk))
 
     if not result: