Optimize queries

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

Author: TG1999

vulnerabilities/api_v3.py

@@ -207,11 +207,13 @@ def get_affected_by_vulnerabilities(self, package):
         """Return a dictionary with advisory as keys and their details, including fixed_by_packages."""
         advisories_qs = AdvisoryV2.objects.latest_affecting_advisories_for_purl(package.package_url)
 
-        advisories = list(advisories_qs[:101])
-        if len(advisories) > 100:
+        advisories_ids = advisories_qs.only("id")
+
+        advisories_ids = list(advisories_ids[:101])
+        if len(advisories_ids) > 100:
             return None
 
-        advisory_by_avid = {adv.avid: adv for adv in advisories}
+        advisory_by_avid = {adv.avid: adv for adv in advisories_qs}
         avids = advisory_by_avid.keys()
 
         impacts = (
@@ -222,7 +224,7 @@ def get_affected_by_vulnerabilities(self, package):
 
         impact_by_avid = {impact.advisory.avid: impact for impact in impacts}
 
-        grouped = group_advisories_by_content(advisories)
+        grouped = group_advisories_by_content(advisories_qs)
 
         result = []
         for entry in grouped.values():
@@ -244,30 +246,17 @@ def get_affected_by_vulnerabilities(self, package):
     def get_fixing_vulnerabilities(self, package):
         advisories_qs = AdvisoryV2.objects.latest_fixed_by_advisories_for_purl(package.package_url)
 
-        advisories = list(advisories_qs[:101])
-        if len(advisories) > 100:
-            return None
-
-        advisory_by_avid = {adv.avid: adv for adv in advisories}
-        avids = advisory_by_avid.keys()
-
-        impacts = (
-            package.fixed_in_impacts.filter(advisory__avid__in=avids)
-            .select_related("advisory")
-            .prefetch_related("fixed_by_packages")
-        )
+        advisories_ids = advisories_qs.only("id")
 
-        impact_by_avid = {impact.advisory.avid: impact for impact in impacts}
+        advisories_ids = list(advisories_ids[:101])
+        if len(advisories_ids) > 100:
+            return None
 
-        grouped = group_advisories_by_content(advisories)
+        grouped = group_advisories_by_content(advisories_qs)
 
         result = []
         for entry in grouped.values():
             primary = entry["primary"]
-            impact = impact_by_avid.get(primary.avid)
-            if not impact:
-                continue
-
             result.append(
                 {
                     "advisory_id": primary.avid,
@@ -301,14 +290,24 @@ def create(self, request, *args, **kwargs):
         approximate = serializer.validated_data["approximate"]
 
         if not purls:
-            vulnerable_purls = (
-                PackageV2.objects.vulnerable()
-                .only("package_url")
-                .distinct()
+            pkg_ids = (
+                PackageV2.objects.vulnerable().values_list("id", flat=True)
+                # .distinct()
+            )
+
+            # vulnerable_purls = (
+            #     PackageV2.objects.vulnerable()
+            #     .only("package_url")
+            #     .values_list("package_url", flat=True)
+            #     .distinct()
+            #     .order_by("package_url")
+            # )
+            query = (
+                PackageV2.objects.filter(id__in=pkg_ids)
                 .values_list("package_url", flat=True)
                 .order_by("package_url")
             )
-            page = self.paginate_queryset(vulnerable_purls)
+            page = self.paginate_queryset(query)
             return self.get_paginated_response(page)
 
         plain_purls = None

vulnerabilities/models.py

@@ -3387,7 +3387,7 @@ def vulnerable(self):
         """
         Return only packages that are vulnerable.
         """
-        return self.filter(affected_in_impacts__isnull=False)
+        return self.filter(id__in=ImpactedPackageAffecting.objects.values("package_id").distinct())
 
     def with_is_vulnerable(self):
         """

vulnerabilities/tests/test_api_v3.py

@@ -53,7 +53,7 @@ def test_packages_post_without_details(self):
     def test_packages_post_with_details(self):
         url = reverse("package-v3-list")
 
-        with self.assertNumQueries(21):
+        with self.assertNumQueries(23):
             response = self.client.post(
                 url,
                 data={
@@ -174,7 +174,7 @@ def setUp(self):
     def test_packages_post_purl_with_many_advisories(self):
         url = reverse("package-v3-list")
 
-        with self.assertNumQueries(11):
+        with self.assertNumQueries(12):
             response = self.client.post(
                 url,
                 data={

vulnerabilities/views.py

@@ -179,7 +179,11 @@ def get_context_data(self, **kwargs):
 
     def get_queryset(self):
         purl = self.kwargs.get("purl")
-        return models.AdvisoryV2.objects.latest_affecting_advisories_for_purl(purl).only("advisory_id", "summary", "url", "date_published").prefetch_related("aliases")
+        return (
+            models.AdvisoryV2.objects.latest_affecting_advisories_for_purl(purl)
+            .only("advisory_id", "summary", "url", "date_published")
+            .prefetch_related("aliases")
+        )
 
 
 class FixingAdvisoriesListView(ListView):
@@ -189,7 +193,11 @@ class FixingAdvisoriesListView(ListView):
 
     def get_queryset(self):
         purl = self.kwargs.get("purl")
-        return models.AdvisoryV2.objects.latest_fixed_by_advisories_for_purl(purl).only("advisory_id", "summary", "url", "date_published").prefetch_related("aliases")
+        return (
+            models.AdvisoryV2.objects.latest_fixed_by_advisories_for_purl(purl)
+            .only("advisory_id", "summary", "url", "date_published")
+            .prefetch_related("aliases")
+        )
 
 
 class PackageV2Details(DetailView):