Use paginated queryset for better memory performance

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

vulnerabilities/pipelines/flag_ghost_packages.py

@@ -8,10 +8,11 @@
 #
 
 import logging
+from itertools import groupby
 from traceback import format_exc as traceback_format_exc
 
 from aboutcode.pipeline import LoopProgress
-from fetchcode.package_versions import SUPPORTED_ECOSYSTEMS
+from fetchcode.package_versions import SUPPORTED_ECOSYSTEMS as FETCHCODE_SUPPORTED_ECOSYSTEMS
 from fetchcode.package_versions import versions
 from packageurl import PackageURL
 from univers.version_range import RANGE_CLASS_BY_SCHEMES
@@ -32,74 +33,76 @@ def flag_ghost_packages(self):
 
 
 def detect_and_flag_ghost_packages(logger=None):
-    """Use fetchcode to validate the package indeed exists upstream."""
+    """Check if packages are available upstream. If not, mark them as ghost package."""
     interesting_packages_qs = (
-        Package.objects.filter(type__in=SUPPORTED_ECOSYSTEMS)
+        Package.objects.order_by("type", "namespace", "name")
+        .filter(type__in=FETCHCODE_SUPPORTED_ECOSYSTEMS)
         .filter(qualifiers="")
         .filter(subpath="")
     )
 
-    distinct_packages = interesting_packages_qs.values("type", "namespace", "name").distinct(
-        "type", "namespace", "name"
+    distinct_packages_count = (
+        interesting_packages_qs.values("type", "namespace", "name")
+        .distinct("type", "namespace", "name")
+        .count()
     )
 
-    distinct_packages_count = distinct_packages.count()
-    package_iterator = distinct_packages.iterator(chunk_size=2000)
-    progress = LoopProgress(total_iterations=distinct_packages_count, logger=logger)
+    grouped_packages = groupby(
+        interesting_packages_qs.paginated(),
+        key=lambda pkg: (pkg.type, pkg.namespace, pkg.name),
+    )
 
     ghost_package_count = 0
-
-    for package in progress.iter(package_iterator):
+    progress = LoopProgress(total_iterations=distinct_packages_count, logger=logger)
+    for type_namespace_name, packages in progress.iter(grouped_packages):
         ghost_package_count += flag_ghost_package(
-            package_dict=package,
-            interesting_packages_qs=interesting_packages_qs,
+            base_purl=PackageURL(*type_namespace_name),
+            packages=packages,
             logger=logger,
         )
 
     if logger:
         logger(f"Successfully flagged {ghost_package_count:,d} ghost Packages")
 
 
-def flag_ghost_package(package_dict, interesting_packages_qs, logger=None):
+def flag_ghost_package(base_purl, packages, logger=None):
     """
-    Check if all the versions of the package described by `package_dict` (type, namespace, name)
-    are available upstream. If they are not available, update the status to 'ghost'.
-    Otherwise, update the status to 'valid'.
+    Check if all the versions of the `purl` are available upstream.
+    If they are not available, update the `is_ghost` to `True`.
     """
-    if not package_dict["type"] in RANGE_CLASS_BY_SCHEMES:
+    if not base_purl.type in RANGE_CLASS_BY_SCHEMES:
         return 0
 
-    known_versions = get_versions(**package_dict, logger=logger)
-    if not known_versions:
+    known_versions = get_versions(purl=base_purl, logger=logger)
+    # Skip if encounter error while fetching known versions
+    if known_versions is None:
         return 0
 
-    version_class = RANGE_CLASS_BY_SCHEMES[package_dict["type"]].version_class
-    package_versions = interesting_packages_qs.filter(**package_dict).filter(status="unknown")
-
     ghost_packages = 0
-    for pkg in package_versions:
+    version_class = RANGE_CLASS_BY_SCHEMES[base_purl.type].version_class
+    for pkg in packages:
+        pkg.is_ghost = False
         if version_class(pkg.version) not in known_versions:
-            pkg.status = "ghost"
-            pkg.save()
+            pkg.is_ghost = True
             ghost_packages += 1
 
-    valid_package_versions = package_versions.exclude(status="ghost")
-    valid_package_versions.update(status="valid")
+            if logger:
+                logger(f"Flagging ghost package {pkg.purl!s}", level=logging.DEBUG)
+        pkg.save()
 
     return ghost_packages
 
 
-def get_versions(type, namespace, name, logger=None):
-    """Return set of known versions for the given package type, namespace, and name."""
-    versionless_purl = PackageURL(type=type, namespace=namespace, name=name)
-    version_class = RANGE_CLASS_BY_SCHEMES[type].version_class
+def get_versions(purl, logger=None):
+    """Return set of known versions for the given purl."""
+    version_class = RANGE_CLASS_BY_SCHEMES[purl.type].version_class
 
     try:
-        return {version_class(v.value) for v in versions(str(versionless_purl))}
+        return {version_class(v.value) for v in versions(str(purl))}
     except Exception as e:
         if logger:
             logger(
-                f"Error while fetching known versions for {versionless_purl!r}: {e!r} \n {traceback_format_exc()}",
+                f"Error while fetching known versions for {purl!s}: {e!r} \n {traceback_format_exc()}",
                 level=logging.ERROR,
             )
         return

vulnerabilities/tests/pipelines/__init__.py

@@ -0,0 +1,20 @@
+#
+# Copyright (c) nexB Inc. and others. All rights reserved.
+# VulnerableCode is a trademark of nexB Inc.
+# SPDX-License-Identifier: Apache-2.0
+# See http://www.apache.org/licenses/LICENSE-2.0 for the license text.
+# See https://github.com/nexB/vulnerablecode for support or download.
+# See https://aboutcode.org for more information about nexB OSS projects.
+#
+
+import io
+
+
+class TestLogger:
+    buffer = io.StringIO()
+
+    def write(self, msg, level=None):
+        self.buffer.write(msg)
+
+    def getvalue(self):
+        return self.buffer.getvalue()

vulnerabilities/tests/pipelines/test_flag_ghost_packages.py

@@ -7,15 +7,17 @@
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
 
-import io
+
 from pathlib import Path
 from unittest import mock
 
 from django.test import TestCase
 from fetchcode.package_versions import PackageVersion
+from packageurl import PackageURL
 
 from vulnerabilities.models import Package
 from vulnerabilities.pipelines import flag_ghost_packages
+from vulnerabilities.tests.pipelines import TestLogger
 
 
 class FlagGhostPackagePipelineTest(TestCase):
@@ -30,20 +32,16 @@ def test_flag_ghost_package(self, mock_fetchcode_versions):
             PackageVersion(value="2.3.0"),
         ]
         interesting_packages_qs = Package.objects.all()
-        target_package = {
-            "type": "pypi",
-            "namespace": "",
-            "name": "foo",
-        }
+        base_purl = PackageURL(type="pypi", name="foo")
 
-        self.assertEqual(0, Package.objects.filter(status="ghost").count())
+        self.assertEqual(0, Package.objects.filter(is_ghost=True).count())
 
         flagged_package_count = flag_ghost_packages.flag_ghost_package(
-            package_dict=target_package,
-            interesting_packages_qs=interesting_packages_qs,
+            base_purl=base_purl,
+            packages=interesting_packages_qs,
         )
         self.assertEqual(1, flagged_package_count)
-        self.assertEqual(1, Package.objects.filter(status="ghost").count())
+        self.assertEqual(1, Package.objects.filter(is_ghost=True).count())
 
     @mock.patch("vulnerabilities.pipelines.flag_ghost_packages.versions")
     def test_detect_and_flag_ghost_packages(self, mock_fetchcode_versions):
@@ -62,11 +60,12 @@ def test_detect_and_flag_ghost_packages(self, mock_fetchcode_versions):
         ]
 
         self.assertEqual(3, Package.objects.count())
-        self.assertEqual(0, Package.objects.filter(status="ghost").count())
+        self.assertEqual(0, Package.objects.filter(is_ghost=True).count())
+
+        logger = TestLogger()
 
-        buffer = io.StringIO()
-        flag_ghost_packages.detect_and_flag_ghost_packages(logger=buffer.write)
+        flag_ghost_packages.detect_and_flag_ghost_packages(logger=logger.write)
         expected = "Successfully flagged 1 ghost Packages"
 
-        self.assertIn(expected, buffer.getvalue())
-        self.assertEqual(1, Package.objects.filter(status="ghost").count())
+        self.assertIn(expected, logger.getvalue())
+        self.assertEqual(1, Package.objects.filter(is_ghost=True).count())