Make improvers query correct and faster

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

Author: TG1999

vulnerabilities/migrations/0117_advisoryv2_risk_score.py

@@ -0,0 +1,24 @@
+# Generated by Django 5.2.11 on 2026-03-17 09:07
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("vulnerabilities", "0116_advisoryv2_advisory_content_hash"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="advisoryv2",
+            name="risk_score",
+            field=models.DecimalField(
+                blank=True,
+                decimal_places=1,
+                help_text="Risk expressed as a number ranging from 0 to 10. Risk is calculated from weighted severity and exploitability values. It is the maximum value of (the weighted severity multiplied by its exploitability) or 10. Risk = min(weighted severity * exploitability, 10)",
+                max_digits=3,
+                null=True,
+            ),
+        ),
+    ]

vulnerabilities/models.py

@@ -3047,17 +3047,13 @@ class AdvisoryV2(models.Model):
         help_text="A unique hash computed from the content of the advisory used to identify advisories with the same content.",
     )
 
-    @property
-    def risk_score(self):
-        """
-        Risk expressed as a number ranging from 0 to 10.
-        Risk is calculated from weighted severity and exploitability values.
-        It is the maximum value of (the weighted severity multiplied by its exploitability) or 10
-        Risk = min(weighted severity * exploitability, 10)
-        """
-        if self.exploitability and self.weighted_severity:
-            risk_score = min(float(self.exploitability * self.weighted_severity), 10.0)
-            return round(risk_score, 1)
+    risk_score = models.DecimalField(
+        null=True,
+        blank=True,
+        max_digits=3,
+        decimal_places=1,
+        help_text="Risk expressed as a number ranging from 0 to 10. Risk is calculated from weighted severity and exploitability values. It is the maximum value of (the weighted severity multiplied by its exploitability) or 10. Risk = min(weighted severity * exploitability, 10)",
+    )
 
     objects = AdvisoryV2QuerySet.as_manager()
 

vulnerabilities/pipelines/v2_improvers/collect_ssvc_trees.py

@@ -36,7 +36,8 @@ def steps(cls):
 
     def collect_ssvc_data(self):
         vulnrichment_advisories = (
-            AdvisoryV2.objects.filter(
+            AdvisoryV2.objects.latest_per_avid()
+            .filter(
                 severities__scoring_system=SCORING_SYSTEMS["ssvc"],
             )
             .distinct()

vulnerabilities/pipelines/v2_improvers/compute_advisory_content_hash.py

@@ -27,7 +27,7 @@ def steps(cls):
     def compute_advisory_content_hash(self):
         """Compute Advisory Content Hash for Advisory."""
 
-        advisories = AdvisoryV2.objects.filter(advisory_content_hash__isnull=True)
+        advisories = AdvisoryV2.objects.latest_per_avid().filter(advisory_content_hash__isnull=True)
 
         advisories_count = advisories.count()
 

vulnerabilities/pipelines/v2_improvers/compute_package_risk.py

@@ -8,8 +8,9 @@
 #
 from aboutcode.pipeline import LoopProgress
 from django.db.models import Prefetch
-from django.db.models import Q
 
+from vulnerabilities.models import AdvisoryExploit
+from vulnerabilities.models import AdvisoryReference
 from vulnerabilities.models import AdvisorySeverity
 from vulnerabilities.models import AdvisoryV2
 from vulnerabilities.models import PackageV2
@@ -36,61 +37,92 @@ def steps(cls):
         )
 
     def compute_and_store_vulnerability_risk_score(self):
+
         affected_advisories = (
-            AdvisoryV2.objects.filter(impacted_packages__affecting_packages__isnull=False)
+            AdvisoryV2.objects.latest_per_avid()
+            .filter(impacted_packages__affecting_packages__isnull=False)
+            .only("id")
             .prefetch_related(
-                "references",
-                "severities",
-                "exploits",
+                Prefetch(
+                    "references", queryset=AdvisoryReference.objects.only("id", "reference_type")
+                ),
+                Prefetch(
+                    "severities",
+                    queryset=AdvisorySeverity.objects.only("id", "value", "url", "scoring_system"),
+                ),
+                Prefetch("exploits", queryset=AdvisoryExploit.objects.only("id")),
                 Prefetch(
                     "related_advisory_severities",
-                    queryset=AdvisoryV2.objects.prefetch_related("severities"),
+                    queryset=AdvisoryV2.objects.only("id").prefetch_related(
+                        Prefetch(
+                            "severities",
+                            queryset=AdvisorySeverity.objects.only(
+                                "id", "value", "url", "scoring_system"
+                            ),
+                        )
+                    ),
                 ),
             )
             .distinct()
         )
 
+        estimated_vulnerability_count = affected_advisories.count()
+
         self.log(
-            f"Calculating risk for {affected_advisories.count():,d} advisory with a affected packages records"
+            f"Calculating risk for {estimated_vulnerability_count:,d} advisory with a affected packages records"
         )
 
-        progress = LoopProgress(total_iterations=affected_advisories.count(), logger=self.log)
+        progress = LoopProgress(
+            logger=self.log, total_iterations=estimated_vulnerability_count, progress_step=5
+        )
 
         updatables = []
         updated_vulnerability_count = 0
         batch_size = 5000
 
         for advisory in progress.iter(affected_advisories.iterator(chunk_size=batch_size)):
+
             references = advisory.references.all()
             exploits = advisory.exploits.all()
 
-            severities = AdvisorySeverity.objects.filter(
-                Q(advisories=advisory) | Q(advisories__related_to_advisory_severities=advisory)
-            ).distinct()
+            severities = list(advisory.severities.all())
+
+            for rel in advisory.related_advisory_severities.all():
+                severities.extend(rel.severities.all())
 
-            weighted_severity, exploitability = compute_vulnerability_risk_factors(
-                references=references,
-                severities=severities,
-                exploits=exploits,
-            )
-            advisory.weighted_severity = weighted_severity
-            advisory.exploitability = exploitability
-            updatables.append(advisory)
+
+            try:
+                weighted_severity, exploitability = compute_vulnerability_risk_factors(
+                    references=references,
+                    severities=severities,
+                    exploits=exploits,
+                )
+
+                advisory.weighted_severity = weighted_severity
+                advisory.exploitability = exploitability
+                if advisory.exploitability and advisory.weighted_severity:
+                    risk_score = min(float(advisory.exploitability * advisory.weighted_severity), 10.0)
+                    advisory.risk_score = round(risk_score, 1)
+                updatables.append(advisory)
+            except Exception as e:
+                self.log(f"Error computing risk score for advisory {advisory.advisory_id}: {e}")
 
             if len(updatables) >= batch_size:
                 updated_vulnerability_count += bulk_update(
                     model=AdvisoryV2,
                     items=updatables,
-                    fields=["weighted_severity", "exploitability"],
+                    fields=["weighted_severity", "exploitability", "risk_score"],
                     logger=self.log,
                 )
-
-        updated_vulnerability_count += bulk_update(
-            model=AdvisoryV2,
-            items=updatables,
-            fields=["weighted_severity", "exploitability"],
-            logger=self.log,
-        )
+                updatables.clear()
+
+        if updatables:
+            updated_vulnerability_count += bulk_update(
+                model=AdvisoryV2,
+                items=updatables,
+                fields=["weighted_severity", "exploitability", "risk_score"],
+                logger=self.log,
+            )
 
         self.log(
             f"Successfully added risk score for {updated_vulnerability_count:,d} vulnerability"
@@ -109,17 +141,19 @@ def compute_and_store_package_risk_score(self):
 
         updatables = []
         updated_package_count = 0
-        batch_size = 10000
+        batch_size = 1000
 
         for package in progress.iter(affected_packages.iterator(chunk_size=batch_size)):
-            risk_score = compute_package_risk_v2(package)
-
-            if not risk_score:
+            try:
+                risk_score = compute_package_risk_v2(package)
+                if not risk_score:
+                    continue
+                package.risk_score = risk_score
+                updatables.append(package)
+            except Exception as e:
+                self.log(f"Error computing risk score for package {package.purl}: {e}")
                 continue
 
-            package.risk_score = risk_score
-            updatables.append(package)
-
             if len(updatables) >= batch_size:
                 updated_package_count += bulk_update(
                     model=PackageV2,

vulnerabilities/pipelines/v2_improvers/enhance_with_exploitdb.py

@@ -89,7 +89,7 @@ def add_vulnerability_exploit(row, logger):
                 for adv in alias.advisories.all():
                     advisories.add(adv)
             else:
-                advs = AdvisoryV2.objects.filter(advisory_id=raw_alias)
+                advs = AdvisoryV2.objects.filter(advisory_id=raw_alias).latest_per_avid()
                 for adv in advs:
                     advisories.add(adv)
         except AdvisoryAlias.DoesNotExist:

vulnerabilities/pipelines/v2_improvers/enhance_with_kev.py

@@ -78,7 +78,7 @@ def add_vulnerability_exploit(kev_vul, logger):
             for adv in alias.advisories.all():
                 advisories.add(adv)
         else:
-            advs = AdvisoryV2.objects.filter(advisory_id=cve_id)
+            advs = AdvisoryV2.objects.filter(advisory_id=cve_id).latest_per_avid()
             for adv in advs:
                 advisories.add(adv)
     except AdvisoryAlias.DoesNotExist:

vulnerabilities/pipelines/v2_improvers/enhance_with_metasploit.py

@@ -83,7 +83,7 @@ def add_advisory_exploit(record, logger):
                 for adv in alias.advisories.all():
                     advisories.add(adv)
             else:
-                advs = AdvisoryV2.objects.filter(advisory_id=ref)
+                advs = AdvisoryV2.objects.filter(advisory_id=ref).latest_per_avid()
                 for adv in advs:
                     advisories.add(adv)
         except AdvisoryAlias.DoesNotExist:

vulnerabilities/pipelines/v2_improvers/relate_severities.py

@@ -61,23 +61,30 @@ def relate_severities(self):
         severity_score_advisories = (
             AdvisoryV2.objects.filter(datasource_id__in=self.pipelines)
             .filter(severities__scoring_system__in=self.SUPPORTED_SYSTEMS)
-            .distinct()
             .latest_per_avid()
+            .distinct()
         )
 
         total = severity_score_advisories.count()
         self.log(f"Processing {total:,d} advisories records")
 
         advisory_id_map = {}
 
-        qs = AdvisoryV2.objects.filter(
-            advisory_id__in=severity_score_advisories.values("advisory_id")
-        ).values("id", "advisory_id")
-
-        alias_qs = AdvisoryV2.objects.filter(
-            aliases__alias__in=severity_score_advisories.values("advisory_id")
-        ).values("id", "aliases__alias")
+        qs = (
+            AdvisoryV2.objects.filter(
+                advisory_id__in=severity_score_advisories.values("advisory_id")
+            )
+            .latest_per_avid()
+            .values("id", "advisory_id")
+        )
 
+        alias_qs = (
+            AdvisoryV2.objects.filter(
+                aliases__alias__in=severity_score_advisories.values("advisory_id")
+            )
+            .latest_per_avid()
+            .values("id", "aliases__alias")
+        )
         for row in qs:
             advisory_id_map.setdefault(row["advisory_id"], set()).add(row["id"])
 

vulnerabilities/risk.py

@@ -8,9 +8,11 @@
 #
 from urllib.parse import urlparse
 
-from vulnerabilities.models import VulnerabilityReference
+from vulnerabilities.models import AdvisoryV2, VulnerabilityReference
 from vulnerabilities.severity_systems import EPSS
 from vulnerabilities.weight_config import WEIGHT_CONFIG
+from django.db.models import Max
+
 
 DEFAULT_WEIGHT = 5
 
@@ -123,12 +125,14 @@ def compute_package_risk_v2(package):
     Calculate the risk for a package by iterating over all vulnerabilities that affects this package
     and determining the associated risk.
     """
-    result = []
-    for impact in package.affected_in_impacts.all():
-        if risk := impact.advisory.risk_score:
-            result.append(float(risk))
 
-    if not result:
+    max_risk = (
+        AdvisoryV2.objects
+        .latest_affecting_advisories_for_purl(package.purl)
+        .aggregate(max_risk=Max("risk_score"))
+    )["max_risk"]
+
+    if max_risk is None:
         return
 
-    return round(max(result), 1)
+    return round(float(max_risk), 1)