Add purl annotation methods on the PackageQuerySet #276

Signed-off-by: tdruez <[email protected]>

Author: tdruez

component_catalog/models.py

@@ -18,10 +18,14 @@
 from django.core.exceptions import ValidationError
 from django.core.validators import EMPTY_VALUES
 from django.db import models
+from django.db.models import Case
 from django.db.models import CharField
 from django.db.models import Count
 from django.db.models import Exists
+from django.db.models import F
 from django.db.models import OuterRef
+from django.db.models import Value
+from django.db.models import When
 from django.db.models.functions import Concat
 from django.dispatch import receiver
 from django.template.defaultfilters import filesizeformat
@@ -1651,6 +1655,65 @@ def __str__(self):
 PACKAGE_URL_FIELDS = ["type", "namespace", "name", "version", "qualifiers", "subpath"]
 
 
+def get_plain_package_url_expression():
+    """
+    Return a Django expression to compute the "PLAIN" Package URL (PURL).
+    Return an empty string if the required `type` or `name` values are missing.
+    """
+    plain_package_url = Concat(
+        Value("pkg:"),
+        F("type"),
+        Case(
+            When(namespace="", then=Value("")),
+            default=Concat(Value("/"), F("namespace")),
+            output_field=CharField(),
+        ),
+        Value("/"),
+        F("name"),
+        Case(
+            When(version="", then=Value("")),
+            default=Concat(Value("@"), F("version")),
+            output_field=CharField(),
+        ),
+        output_field=CharField(),
+    )
+
+    return Case(
+        When(type="", then=Value("")),
+        When(name="", then=Value("")),
+        default=plain_package_url,
+        output_field=CharField(),
+    )
+
+
+def get_package_url_expression():
+    """
+    Return a Django expression to compute the "FULL" Package URL (PURL).
+    Return an empty string if the required `type` or `name` values are missing.
+    """
+    package_url = Concat(
+        get_plain_package_url_expression(),
+        Case(
+            When(qualifiers="", then=Value("")),
+            default=Concat(Value("?"), F("qualifiers")),
+            output_field=CharField(),
+        ),
+        Case(
+            When(subpath="", then=Value("")),
+            default=Concat(Value("#"), F("subpath")),
+            output_field=CharField(),
+        ),
+        output_field=CharField(),
+    )
+
+    return Case(
+        When(type="", then=Value("")),
+        When(name="", then=Value("")),
+        default=package_url,
+        output_field=CharField(),
+    )
+
+
 class PackageQuerySet(PackageURLQuerySetMixin, VulnerabilityQuerySetMixin, DataspacedQuerySet):
     def has_package_url(self):
         """Return objects with Package URL defined."""
@@ -1666,6 +1729,26 @@ def annotate_sortable_identifier(self):
             sortable_identifier=Concat(*PACKAGE_URL_FIELDS, "filename", output_field=CharField())
         )
 
+    def annotate_plain_package_url(self):
+        """
+        Annotate the QuerySet with a computed 'plain' Package URL (PURL).
+
+        This plain PURL is a simplified version that includes only the core fields:
+        `type`, `namespace`, `name`, and `version`. It omits any qualifiers or
+        subpath components, providing a normalized and minimal representation
+        of the Package URL.
+        """
+        return self.annotate(plain_purl=get_plain_package_url_expression())
+
+    def annotate_package_url(self):
+        """
+        Annotate the QuerySet with a fully-computed Package URL (PURL).
+
+        This includes the core PURL fields (`type`, `namespace`, `name`, `version`)
+        as well as any qualifiers and subpath components.
+        """
+        return self.annotate(purl=get_package_url_expression())
+
     def only_rendering_fields(self):
         """Minimum requirements to render a Package element in the UI."""
         return self.only(

component_catalog/tests/test_models.py

@@ -2709,3 +2709,32 @@ def test_vulnerability_mixin_is_vulnerable_property(self):
         package2 = make_package(self.dataspace)
         self.assertTrue(package1.is_vulnerable)
         self.assertFalse(package2.is_vulnerable)
+
+    def test_package_queryset_has_package_url(self):
+        package1 = make_package(self.dataspace, package_url="pkg:pypi/[email protected]")
+        make_package(self.dataspace)
+        qs = Package.objects.has_package_url()
+        self.assertQuerySetEqual(qs, [package1])
+
+    def test_package_queryset_annotate_sortable_identifier(self):
+        package1 = make_package(self.dataspace, package_url="pkg:pypi/[email protected]")
+        package2 = make_package(self.dataspace)
+        qs = Package.objects.annotate_sortable_identifier()
+        self.assertEqual("pypidjango5.0", qs.get(pk=package1.pk).sortable_identifier)
+        self.assertEqual(package2.filename, qs.get(pk=package2.pk).sortable_identifier)
+
+    def test_package_queryset_annotate_package_url(self):
+        package_url = "pkg:pypi/[email protected]?qualifier=true#path"
+        package1 = make_package(self.dataspace, package_url=package_url)
+        package2 = make_package(self.dataspace)
+        qs = Package.objects.annotate_package_url()
+        self.assertEqual(package_url, qs.get(pk=package1.pk).purl)
+        self.assertEqual("", qs.get(pk=package2.pk).purl)
+
+    def test_package_queryset_annotate_plain_package_url(self):
+        package_url = "pkg:pypi/[email protected]?qualifier=true#path"
+        package1 = make_package(self.dataspace, package_url=package_url)
+        package2 = make_package(self.dataspace)
+        qs = Package.objects.annotate_plain_package_url()
+        self.assertEqual("pkg:pypi/[email protected]", qs.get(pk=package1.pk).plain_purl)
+        self.assertEqual("", qs.get(pk=package2.pk).plain_purl)