Make bulk search fast

Signed-off-by: Tushar Goel <[email protected]>

Author: TG1999

vulnerabilities/api.py

@@ -238,32 +238,26 @@ def bulk_search(self, request):
         """
         Lookup for vulnerable packages using many Package URLs at once.
         """
-        response = []
+
         purls = request.data.get("purls", []) or []
+        purl_only = request.data.get("purl_only", False)
         if not purls or not isinstance(purls, list):
             return Response(
                 status=400,
-                data={"Error": "A non-empty 'purls' list of package URLs is required."},
+                data={"Error": "A non-empty 'purls' list of PURLs is required."},
             )
-        for purl in request.data["purls"]:
-            try:
-                purl_string = purl
-                purl = PackageURL.from_string(purl)
-            except ValueError:
-                return Response(status=400, data={"Error": f"Invalid Package URL: {purl}"})
-            lookups = get_purl_query_lookups(purl)
-            purl_data = Package.objects.filter(**lookups)
-            purl_response = {}
-            if purl_data:
-                purl_response = PackageSerializer(purl_data[0], context={"request": request}).data
-            else:
-                purl_response = purl.to_dict()
-                purl_response["unresolved_vulnerabilities"] = []
-                purl_response["resolved_vulnerabilities"] = []
-                purl_response["purl"] = purl_string
-            response.append(purl_response)
-
-        return Response(response)
+
+        query = Package.objects.filter(package__in=purls)
+
+        if not purl_only:
+            return Response(
+                PackageSerializer(query.distinct(), many=True, context={"request": request}).data
+            )
+
+        vulnerable_purls = (
+            query.filter(packagerelatedvulnerability__fix=False).only("package").distinct()
+        )
+        return Response(data=vulnerable_purls)
 
     @action(detail=False, methods=["get"], throttle_scope="vulnerable_packages")
     def all(self, request):

vulnerabilities/migrations/0034_package_package.py

@@ -0,0 +1,18 @@
+# Generated by Django 4.0.7 on 2022-11-25 12:57
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('vulnerabilities', '0033_alter_vulnerabilityseverity_scoring_system'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='package',
+            name='package',
+            field=models.CharField(blank=True, help_text='The Package URL for this package.', max_length=255),
+        ),
+    ]

vulnerabilities/migrations/0035_auto_20221125_1257.py

@@ -0,0 +1,28 @@
+# Generated by Django 4.0.7 on 2022-11-25 12:57
+
+from django.db import migrations
+from packageurl import PackageURL
+
+class Migration(migrations.Migration):
+
+    def save_package(apps, schema_editor):
+        Package = apps.get_model("vulnerabilities", "Package")
+        for package in Package.objects.all():
+            purl = PackageURL(
+                type=package.type,
+                namespace=package.namespace,
+                name=package.name,
+                version=package.version,
+                qualifiers=package.qualifiers,
+                subpath=package.subpath,
+            )
+            package.package = str(purl)
+            package.save()
+
+    dependencies = [
+        ('vulnerabilities', '0034_package_package'),
+    ]
+
+    operations = [
+        migrations.RunPython(save_package, reverse_code=migrations.RunPython.noop),
+    ]

vulnerabilities/models.py

@@ -531,8 +531,19 @@ class Package(PackageURLMixin):
         to="Vulnerability", through="PackageRelatedVulnerability"
     )
 
+    package = models.CharField(
+        max_length=255,
+        blank=True,
+        null=False,
+        help_text="The Package URL for this package.",
+    )
+
     objects = PackageQuerySet.as_manager()
 
+    def save(self, *args, **kwargs):
+        self.package = self.purl
+        super().save(*args, **kwargs)
+
     @property
     def purl(self):
         return self.package_url