Optimize V3 API

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

Author: TG1999

vulnerabilities/api_v3.py

@@ -11,14 +11,11 @@
 from urllib.parse import urlencode
 
 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
 from django.db.models import OuterRef
 from django.db.models import Prefetch
 from django.db.models import Q
-from django.db.models.functions import JSONObject
 from django_filters import rest_framework as filters
 from drf_spectacular.utils import extend_schema
 from packageurl import PackageURL
@@ -31,6 +28,7 @@
 from vulnerabilities.models import AdvisoryAlias
 from vulnerabilities.models import AdvisoryReference
 from vulnerabilities.models import AdvisorySet
+from vulnerabilities.models import AdvisorySetMember
 from vulnerabilities.models import AdvisorySeverity
 from vulnerabilities.models import AdvisoryV2
 from vulnerabilities.models import AdvisoryWeakness
@@ -503,7 +501,10 @@ def get_affected_advisories_bulk(packages, max_advisories, base_url):
         )
         .select_related("primary_advisory")
         .prefetch_related(
-            Prefetch("aliases", queryset=AdvisoryAlias.objects.only("alias")),
+            Prefetch(
+                "aliases",
+                queryset=AdvisoryAlias.objects.only("alias"),
+            ),
         )
         .annotate(
             max_severity=Max(
@@ -512,16 +513,6 @@ def get_affected_advisories_bulk(packages, max_advisories, base_url):
             max_exploitability=Max(
                 "members__advisory__exploitability",
             ),
-            ssvc_trees=JSONBAgg(
-                JSONObject(
-                    vector=F("members__advisory__related_ssvcs__vector"),
-                    decision=F("members__advisory__related_ssvcs__decision"),
-                    options=F("members__advisory__related_ssvcs__options"),
-                    source_url=F("members__advisory__related_ssvcs__source_advisory__url"),
-                ),
-                filter=Q(members__advisory__related_ssvcs__decision__isnull=False),
-                distinct=True,
-            ),
         )
         .only(
             "id",
@@ -532,10 +523,88 @@ def get_affected_advisories_bulk(packages, max_advisories, base_url):
         )
     )
 
+    advisory_set_ids = [a.id for a in advisory_sets]
+
+    member_rows = AdvisorySetMember.objects.filter(
+        advisory_set_id__in=advisory_set_ids,
+    ).values(
+        "advisory_set_id",
+        "advisory_id",
+    )
+
+    advisory_ids_by_set = defaultdict(set)
+    all_advisory_ids = set()
+
+    for row in member_rows:
+        advisory_set_id = row["advisory_set_id"]
+        advisory_id = row["advisory_id"]
+
+        advisory_ids_by_set[advisory_set_id].add(advisory_id)
+        all_advisory_ids.add(advisory_id)
+
+    ssvc_rows = (
+        SSVC.objects.filter(
+            related_advisories__id__in=all_advisory_ids,
+            decision__isnull=False,
+        )
+        .select_related(
+            "source_advisory",
+        )
+        .values(
+            "related_advisories__id",
+            "vector",
+            "decision",
+            "options",
+            "source_advisory__url",
+        )
+    )
+
+    ssvc_by_advisory = defaultdict(list)
+
+    for row in ssvc_rows:
+        advisory_id = row["related_advisories__id"]
+
+        ssvc_by_advisory[advisory_id].append(
+            {
+                "vector": row["vector"],
+                "decision": row["decision"],
+                "options": row["options"],
+                "source_url": row["source_advisory__url"],
+            }
+        )
+
     package_map = defaultdict(list)
 
     for adv in advisory_sets:
         adv._aliases_cache = [a.alias for a in adv.aliases.all()]
+
+        advisory_ids = advisory_ids_by_set.get(adv.id, set())
+
+        seen = set()
+        ssvc_trees = []
+
+        for advisory_id in advisory_ids:
+            for ssvc in ssvc_by_advisory.get(advisory_id, []):
+
+                key = (
+                    ssvc["vector"],
+                    ssvc["decision"],
+                    (
+                        tuple(sorted(ssvc["options"].items()))
+                        if isinstance(ssvc["options"], dict)
+                        else str(ssvc["options"])
+                    ),
+                    ssvc["source_url"],
+                )
+
+                if key in seen:
+                    continue
+
+                seen.add(key)
+                ssvc_trees.append(ssvc)
+
+        adv.ssvc_trees = ssvc_trees
+
         package_map[adv.package_id].append(adv)
 
     for package in packages:
@@ -549,6 +618,7 @@ def get_affected_advisories_bulk(packages, max_advisories, base_url):
                 (package.id, primary.id),
                 [],
             )
+
             max_sev = adv.max_severity or 0.0
             max_exp = adv.max_exploitability or 0.0
 
@@ -577,7 +647,7 @@ def get_affected_advisories_bulk(packages, max_advisories, base_url):
                     "exploitability": exploitability,
                     "risk_score": risk_score,
                     "fixed_by_packages": fixed_by_packages,
-                    "ssvc_trees": adv.ssvc_trees or [],
+                    "ssvc_trees": adv.ssvc_trees,
                     "resource_url": resource_url,
                 }
             )

vulnerabilities/pipelines/v2_improvers/mark_unfurl_version_range.py

@@ -50,7 +50,7 @@ def mark_all_impacts_unfurled(self):
             impacted_packages=impacted_packages
         )
 
-        batch_size = 100
+        batch_size = 1000
         batch = []
 
         successful_qs = latest_advisories_with_all_impacts_unfurled_successfully(

vulnerabilities/tests/pipelines/v2_improvers/test_mark_all_impacts_unfurled.py

@@ -194,15 +194,10 @@ def test_attempted_advisories_are_chunked_in_batches_of_100(
 
         pipeline.mark_all_impacts_unfurled()
 
-        assert mock_complete_advisories_import.call_count == 3
+        assert mock_complete_advisories_import.call_count == 1
 
         first_call_ids = mock_complete_advisories_import.call_args_list[0][1]["advisory_ids"]
-        second_call_ids = mock_complete_advisories_import.call_args_list[1][1]["advisory_ids"]
-        third_call_ids = mock_complete_advisories_import.call_args_list[2][1]["advisory_ids"]
-
-        assert len(first_call_ids) == 100
-        assert len(second_call_ids) == 100
-        assert len(third_call_ids) == 50
+        assert len(first_call_ids) == 250
 
 
 @pytest.mark.django_db