Merge pull request #1558 from aboutcode-org/add-tests-for-queries

Enhance API performace

Author: TG1999

vulnerabilities/api.py

@@ -87,29 +87,23 @@ def get_resource_url(self, instance):
         return resource_url
 
 
-class MinimalPackageSerializer(BaseResourceSerializer):
+class VulnVulnIDSerializer(serializers.Serializer):
     """
-    Used for nesting inside vulnerability focused APIs.
+    Serializer for the series of vulnerability IDs.
     """
 
-    def get_affected_vulnerabilities(self, package):
-        parent_affected_vulnerabilities = package.fixed_package_details.get("vulnerabilities") or []
-
-        affected_vulnerabilities = [
-            self.get_vulnerability(vuln) for vuln in parent_affected_vulnerabilities
-        ]
+    vulnerability = serializers.CharField(source="vulnerability_id")
 
-        return affected_vulnerabilities
+    class Meta:
+        fields = ["vulnerability"]
 
-    def get_vulnerability(self, vuln):
-        affected_vulnerability = {}
 
-        vulnerability = vuln.get("vulnerability")
-        if vulnerability:
-            affected_vulnerability["vulnerability"] = vulnerability.vulnerability_id
-            return affected_vulnerability
+class MinimalPackageSerializer(BaseResourceSerializer):
+    """
+    Used for nesting inside vulnerability focused APIs.
+    """
 
-    affected_by_vulnerabilities = serializers.SerializerMethodField("get_affected_vulnerabilities")
+    affected_by_vulnerabilities = VulnVulnIDSerializer(source="affecting_vulns", many=True)
 
     purl = serializers.CharField(source="package_url")
 
@@ -145,18 +139,17 @@ class VulnSerializerRefsAndSummary(BaseResourceSerializer):
     Lookup vulnerabilities references by aliases (such as a CVE).
     """
 
-    def to_representation(self, instance):
-        data = super().to_representation(instance)
-        aliases = [alias["alias"] for alias in data["aliases"]]
-        data["aliases"] = aliases
-        return data
-
     fixed_packages = MinimalPackageSerializer(
         many=True, source="filtered_fixed_packages", read_only=True
     )
 
     references = VulnerabilityReferenceSerializer(many=True, source="vulnerabilityreference_set")
-    aliases = AliasSerializer(many=True, source="alias")
+
+    aliases = serializers.SerializerMethodField()
+
+    def get_aliases(self, obj):
+        # Assuming `obj.aliases` is a queryset of `Alias` objects
+        return [alias.alias for alias in obj.aliases.all()]
 
     class Meta:
         model = Vulnerability
@@ -257,34 +250,22 @@ class PackageSerializer(BaseResourceSerializer):
     Lookup software package using Package URLs
     """
 
-    def to_representation(self, instance):
-        data = super().to_representation(instance)
-        data["qualifiers"] = normalize_qualifiers(data["qualifiers"], encode=False)
-
-        return data
-
-    next_non_vulnerable_version = serializers.SerializerMethodField("get_next_non_vulnerable")
-
-    def get_next_non_vulnerable(self, package):
-        next_non_vulnerable = package.fixed_package_details.get("next_non_vulnerable", None)
-        if next_non_vulnerable:
-            return next_non_vulnerable.version
-
-    latest_non_vulnerable_version = serializers.SerializerMethodField("get_latest_non_vulnerable")
-
-    def get_latest_non_vulnerable(self, package):
-        latest_non_vulnerable = package.fixed_package_details.get("latest_non_vulnerable", None)
-        if latest_non_vulnerable:
-            return latest_non_vulnerable.version
+    next_non_vulnerable_version = serializers.CharField(read_only=True)
+    latest_non_vulnerable_version = serializers.CharField(read_only=True)
 
     purl = serializers.CharField(source="package_url")
 
     affected_by_vulnerabilities = serializers.SerializerMethodField("get_affected_vulnerabilities")
 
     fixing_vulnerabilities = serializers.SerializerMethodField("get_fixing_vulnerabilities")
 
+    qualifiers = serializers.SerializerMethodField()
+
     is_vulnerable = serializers.BooleanField()
 
+    def get_qualifiers(self, package):
+        return normalize_qualifiers(package.qualifiers, encode=False)
+
     def get_fixed_packages(self, package):
         """
         Return a queryset of all packages that fix a vulnerability with
@@ -368,8 +349,6 @@ class Meta:
             "fixing_vulnerabilities",
         ]
 
-    is_vulnerable = serializers.BooleanField()
-
 
 class PackageFilterSet(filters.FilterSet):
     purl = filters.CharFilter(method="filter_purl")

vulnerabilities/models.py

@@ -711,6 +711,47 @@ def version_class(self):
     def current_version(self):
         return self.version_class(self.version)
 
+    @property
+    def next_non_vulnerable_version(self):
+        """
+        Return the version string of the next non-vulnerable package version.
+        """
+        next_non_vulnerable, _ = self.get_non_vulnerable_versions()
+        return next_non_vulnerable.version if next_non_vulnerable else None
+
+    @property
+    def latest_non_vulnerable_version(self):
+        """
+        Return the version string of the latest non-vulnerable package version.
+        """
+        _, latest_non_vulnerable = self.get_non_vulnerable_versions()
+        return latest_non_vulnerable.version if latest_non_vulnerable else None
+
+    def get_non_vulnerable_versions(self):
+        """
+        Return a tuple of the next and latest non-vulnerable versions as PackageURL objects.
+        Return a tuple of (None, None) if there is no non-vulnerable version.
+        """
+        non_vulnerable_versions = Package.objects.get_fixed_by_package_versions(
+            self, fix=False
+        ).only_non_vulnerable()
+        sorted_versions = self.sort_by_version(non_vulnerable_versions)
+
+        later_non_vulnerable_versions = [
+            non_vuln_ver
+            for non_vuln_ver in sorted_versions
+            if self.version_class(non_vuln_ver.version) > self.current_version
+        ]
+
+        if later_non_vulnerable_versions:
+            sorted_versions = self.sort_by_version(later_non_vulnerable_versions)
+            next_non_vulnerable_version = sorted_versions[0]
+            latest_non_vulnerable_version = sorted_versions[-1]
+
+            return next_non_vulnerable_version, latest_non_vulnerable_version
+
+        return None, None
+
     @property
     def fixed_package_details(self):
         """
@@ -827,6 +868,20 @@ def affecting_vulnerabilities(self):
         """
         return self.vulnerabilities.filter(packagerelatedvulnerability__fix=False)
 
+    @property
+    def affecting_vulns(self):
+        """
+        Return a queryset of Vulnerabilities that affect this `package`.
+        """
+        fixed_by_packages = Package.objects.get_fixed_by_package_versions(self, fix=True)
+        return self.vulnerabilities.affecting_vulnerabilities().prefetch_related(
+            Prefetch(
+                "packages",
+                queryset=fixed_by_packages,
+                to_attr="fixed_packages",
+            )
+        )
+
 
 class PackageRelatedVulnerability(models.Model):
     """

vulnerabilities/tests/test_api.py

@@ -388,7 +388,7 @@ def add_aliases(vuln, aliases):
         Alias.objects.create(alias=alias, vulnerability=vuln)
 
 
-class APITestCasePackage(TestCase):
+class APIPerformanceTest(TestCase):
     def setUp(self):
         self.user = ApiUser.objects.create_api_user(username="[email protected]")
         self.auth = f"Token {self.user.auth_token.key}"
@@ -439,18 +439,124 @@ def setUp(self):
         set_as_affected_by(package=self.pkg_2_13_2, vulnerability=self.vul2)
         set_as_fixing(package=self.pkg_2_13_2, vulnerability=self.vul1)
 
-    def test_api_with_package_with_no_vulnerabilities(self):
-        affected_vulnerabilities = []
-        vuln = {
-            "foo": "bar",
-        }
+    def test_api_packages_all_num_queries(self):
+        with self.assertNumQueries(4):
+            # There are 4 queries:
+            # 1. SAVEPOINT
+            # 2. Authenticating user
+            # 3. Get all vulnerable packages
+            # 4. RELEASE SAVEPOINT
+            response = self.csrf_client.get(f"/api/packages/all", format="json").data
+
+            assert len(response) == 3
+            assert response == [
+                "pkg:maven/com.fasterxml.jackson.core/[email protected]",
+                "pkg:maven/com.fasterxml.jackson.core/[email protected]",
+                "pkg:maven/com.fasterxml.jackson.core/[email protected]",
+            ]
+
+    def test_api_packages_single_num_queries(self):
+        with self.assertNumQueries(8):
+            self.csrf_client.get(f"/api/packages/{self.pkg_2_14_0_rc1.id}", format="json")
+
+    def test_api_packages_single_with_purl_in_query_num_queries(self):
+        with self.assertNumQueries(9):
+            self.csrf_client.get(f"/api/packages/?purl={self.pkg_2_14_0_rc1.purl}", format="json")
+
+    def test_api_packages_single_with_purl_no_version_in_query_num_queries(self):
+        with self.assertNumQueries(64):
+            self.csrf_client.get(
+                f"/api/packages/?purl=pkg:maven/com.fasterxml.jackson.core/jackson-databind",
+                format="json",
+            )
 
-        package_with_no_vulnerabilities = MinimalPackageSerializer.get_vulnerability(
-            self,
-            vuln,
+    def test_api_packages_bulk_search(self):
+        with self.assertNumQueries(45):
+            packages = [self.pkg_2_12_6, self.pkg_2_12_6_1, self.pkg_2_13_1]
+            purls = [p.purl for p in packages]
+
+            data = {"purls": purls, "purl_only": False, "plain_purl": True}
+
+            resp = self.csrf_client.post(
+                f"/api/packages/bulk_search",
+                data=json.dumps(data),
+                content_type="application/json",
+            ).json()
+
+    def test_api_packages_with_lookup(self):
+        with self.assertNumQueries(14):
+            data = {"purl": self.pkg_2_12_6.purl}
+
+            resp = self.csrf_client.post(
+                f"/api/packages/lookup",
+                data=json.dumps(data),
+                content_type="application/json",
+            ).json()
+
+    def test_api_packages_bulk_lookup(self):
+        with self.assertNumQueries(45):
+            packages = [self.pkg_2_12_6, self.pkg_2_12_6_1, self.pkg_2_13_1]
+            purls = [p.purl for p in packages]
+
+            data = {"purls": purls}
+
+            resp = self.csrf_client.post(
+                f"/api/packages/bulk_lookup",
+                data=json.dumps(data),
+                content_type="application/json",
+            ).json()
+
+
+class APITestCasePackage(TestCase):
+    def setUp(self):
+        self.user = ApiUser.objects.create_api_user(username="[email protected]")
+        self.auth = f"Token {self.user.auth_token.key}"
+        self.csrf_client = APIClient(enforce_csrf_checks=True)
+        self.csrf_client.credentials(HTTP_AUTHORIZATION=self.auth)
+
+        # This setup creates the following data:
+        # vulnerabilities: vul1, vul2, vul3
+        # pkg:maven/com.fasterxml.jackson.core/jackson-databind
+        # with these versions:
+        # pkg_2_12_6:     @ 2.12.6       affected by        fixing vul3
+        # pkg_2_12_6_1:   @ 2.12.6.1     affected by vul2   fixing vul1
+        # pkg_2_13_1:     @ 2.13.1       affected by vul1   fixing vul3
+        # pkg_2_13_2:     @ 2.13.2       affected by vul2   fixing vul1
+        # pkg_2_14_0_rc1: @ 2.14.0-rc1   affected by        fixing
+
+        # searched-for pkg's vuln
+        self.vul1 = create_vuln("VCID-vul1-vul1-vul1", ["CVE-2020-36518", "GHSA-57j2-w4cx-62h2"])
+        self.vul2 = create_vuln("VCID-vul2-vul2-vul2")
+        # This is the vuln fixed by the searched-for pkg -- and by a lesser version (created below),
+        # which WILL be included in the API
+        self.vul3 = create_vuln("VCID-vul3-vul3-vul3", ["CVE-2021-46877", "GHSA-3x8x-79m2-3w2w"])
+
+        from_purl = Package.objects.from_purl
+        # lesser-version pkg that also fixes the vuln fixed by the searched-for pkg
+        self.pkg_2_12_6 = from_purl("pkg:maven/com.fasterxml.jackson.core/[email protected]")
+        # this is a lesser version omitted from the API that fixes searched-for pkg's vuln
+        self.pkg_2_12_6_1 = from_purl(
+            "pkg:maven/com.fasterxml.jackson.core/[email protected]"
+        )
+        # searched-for pkg
+        self.pkg_2_13_1 = from_purl("pkg:maven/com.fasterxml.jackson.core/[email protected]")
+        # this is a greater version that fixes searched-for pkg's vuln
+        self.pkg_2_13_2 = from_purl("pkg:maven/com.fasterxml.jackson.core/[email protected]")
+        # This addresses both next and latest non-vulnerable pkg
+        self.pkg_2_14_0_rc1 = from_purl(
+            "pkg:maven/com.fasterxml.jackson.core/[email protected]"
         )
 
-        assert package_with_no_vulnerabilities is None
+        set_as_fixing(package=self.pkg_2_12_6, vulnerability=self.vul3)
+
+        set_as_affected_by(package=self.pkg_2_12_6_1, vulnerability=self.vul2)
+        set_as_fixing(package=self.pkg_2_12_6_1, vulnerability=self.vul1)
+
+        set_as_affected_by(package=self.pkg_2_13_1, vulnerability=self.vul1)
+        set_as_fixing(package=self.pkg_2_13_1, vulnerability=self.vul3)
+
+        set_as_affected_by(package=self.pkg_2_13_2, vulnerability=self.vul2)
+        set_as_fixing(package=self.pkg_2_13_2, vulnerability=self.vul1)
 
     def test_api_with_lesser_and_greater_fixed_by_packages(self):
         response = self.csrf_client.get(f"/api/packages/{self.pkg_2_13_1.id}", format="json").data

vulnerablecode/settings.py

@@ -287,6 +287,9 @@
 
 
 if DEBUG_TOOLBAR:
+    # Uncomment this to get pyinstrument profiles
+    # PYINSTRUMENT_PROFILE_DIR = "profiles"
+
     INSTALLED_APPS += ("debug_toolbar",)
 
     MIDDLEWARE += (