Speed up SSVC tree collection

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

Author: TG1999

vulnerabilities/pipelines/v2_improvers/collect_ssvc_trees.py

@@ -8,6 +8,7 @@
 #
 
 import logging
+from collections import defaultdict
 
 from django.db.models import Prefetch
 from django.db.models import Q
@@ -35,53 +36,126 @@ def steps(cls):
         return (cls.collect_ssvc_data,)
 
     def collect_ssvc_data(self):
-        vulnrichment_advisories = (
+        advisories = list(
             AdvisoryV2.objects.latest_per_avid()
             .filter(
                 severities__scoring_system=SCORING_SYSTEMS["ssvc"],
             )
-            .distinct()
             .prefetch_related(
                 Prefetch(
                     "severities",
                     queryset=AdvisorySeverity.objects.filter(
                         scoring_system=SCORING_SYSTEMS["ssvc"]
-                    ),
-                )
+                    ).only("id", "scoring_elements"),
+                ),
+                "aliases",
             )
+            .only("id", "advisory_id")
+            .distinct()
         )
 
-        self.log(
-            f"Found {vulnrichment_advisories.count()} advisories from Vulnrichment with SSVC severities."
+        self.log(f"Found {len(advisories)} advisories from Vulnrichment with SSVC severities.")
+        advisory_ids = {a.advisory_id for a in advisories}
+
+        all_related = (
+            AdvisoryV2.objects.filter(
+                Q(advisory_id__in=advisory_ids) | Q(aliases__alias__in=advisory_ids)
+            )
+            .distinct()
+            .only("id", "advisory_id")
+            .prefetch_related("aliases")
         )
-        for advisory in vulnrichment_advisories:
+
+        advisory_map = defaultdict(set)
+
+        for adv in all_related:
+            advisory_map[adv.advisory_id].add(adv)
+            for alias in adv.aliases.all():
+                if alias.alias in advisory_ids:
+                    advisory_map[alias.alias].add(adv)
+
+        existing_ssvc = {
+            s.source_advisory_id: s
+            for s in SSVC.objects.filter(source_advisory_id__in=[a.id for a in advisories])
+        }
+
+        to_create = []
+        to_update = []
+
+        for advisory in advisories:
             self.log(f"Processing advisory: {advisory.advisory_id}")
+
             for severity in advisory.severities.all():
                 ssvc_vector = severity.scoring_elements
-                self.log(f"SSVC Vector found: {ssvc_vector}")
+
                 try:
                     ssvc_tree, decision = convert_vector_to_tree_and_decision(ssvc_vector)
-                    self.log(
-                        f"Advisory: {advisory.advisory_id}, SSVC Tree: {ssvc_tree}, Decision: {decision}, vector: {ssvc_vector}"
-                    )
-                    ssvc_obj, _ = SSVC.objects.get_or_create(
-                        source_advisory=advisory,
-                        defaults={
-                            "options": ssvc_tree,
-                            "decision": decision,
-                            "vector": ssvc_vector,
-                        },
-                    )
-                    # All advisories that have advisory.advisory_id in their aliases or advisory_id same as advisory.advisory_id
-                    related_advisories = AdvisoryV2.objects.filter(
-                        Q(advisory_id=advisory.advisory_id) | Q(aliases__alias=advisory.advisory_id)
-                    ).distinct()
-                    related_advisories = related_advisories.exclude(id=advisory.id)
-                    ssvc_obj.related_advisories.set(related_advisories)
+
+                    existing = existing_ssvc.get(advisory.id)
+
+                    if existing:
+                        existing.options = ssvc_tree
+                        existing.decision = decision
+                        existing.vector = ssvc_vector
+                        to_update.append(existing)
+                        ssvc_obj = existing
+                    else:
+                        ssvc_obj = SSVC(
+                            source_advisory=advisory,
+                            options=ssvc_tree,
+                            decision=decision,
+                            vector=ssvc_vector,
+                        )
+                        to_create.append(ssvc_obj)
+
                 except Exception as e:
                     logger.error(
-                        f"Failed to parse SSVC vector '{ssvc_vector}' for advisory '{advisory}': {e}"
+                        f"Failed to parse SSVC vector '{ssvc_vector}' "
+                        f"for advisory '{advisory}': {e}"
+                    )
+
+        SSVC.objects.bulk_create(to_create, batch_size=1000)
+
+        SSVC.objects.bulk_update(
+            to_update,
+            ["options", "decision", "vector"],
+            batch_size=1000,
+        )
+
+        # Refresh newly created IDs
+        created_ssvc = {
+            s.source_advisory_id: s
+            for s in SSVC.objects.filter(source_advisory_id__in=[a.id for a in advisories])
+        }
+
+        through_model = SSVC.related_advisories.through
+
+        through_rows = []
+
+        for advisory in advisories:
+            ssvc_obj = created_ssvc.get(advisory.id)
+
+            if not ssvc_obj:
+                continue
+
+            related = advisory_map.get(advisory.advisory_id, set())
+
+            for related_adv in related:
+                if related_adv.id == advisory.id:
+                    continue
+
+                through_rows.append(
+                    through_model(
+                        ssvc_id=ssvc_obj.id,
+                        advisoryv2_id=related_adv.id,
                     )
+                )
+
+        through_model.objects.bulk_create(
+            through_rows,
+            ignore_conflicts=True,
+            batch_size=5000,
+        )
 
 
 REVERSE_POINTS = {