Fast content ID migration

vulnerabilities/importer.py

@@ -9,6 +9,7 @@
 
 import dataclasses
 import datetime
+import functools
 import logging
 import os
 import shutil
@@ -46,7 +47,8 @@
 logger = logging.getLogger(__name__)
 
 
-@dataclasses.dataclass(order=True)
+@dataclasses.dataclass(eq=True)
+@functools.total_ordering
 class VulnerabilitySeverity:
     # FIXME: this should be named scoring_system, like in the model
     system: ScoringSystem
@@ -55,15 +57,26 @@ class VulnerabilitySeverity:
     published_at: Optional[datetime.datetime] = None
 
     def to_dict(self):
-        published_at_dict = (
-            {"published_at": self.published_at.isoformat()} if self.published_at else {}
-        )
-        return {
+        data = {
             "system": self.system.identifier,
             "value": self.value,
             "scoring_elements": self.scoring_elements,
-            **published_at_dict,
         }
+        if self.published_at:
+            if isinstance(self.published_at, datetime.datetime):
+                data["published_at"] = self.published_at.isoformat()
+            else:
+                data["published_at"] = self.published_at
+        return data
+
+    def __lt__(self, other):
+        if not isinstance(other, VulnerabilitySeverity):
+            return NotImplemented
+        return self._cmp_key() < other._cmp_key()
+
+    # TODO: Add cache
+    def _cmp_key(self):
+        return (self.system.identifier, self.value, self.scoring_elements, self.published_at)
 
     @classmethod
     def from_dict(cls, severity: dict):
@@ -79,7 +92,8 @@ def from_dict(cls, severity: dict):
         )
 
 
-@dataclasses.dataclass(order=True)
+@dataclasses.dataclass(eq=True)
+@functools.total_ordering
 class Reference:
     reference_id: str = ""
     reference_type: str = ""
@@ -90,27 +104,28 @@ def __post_init__(self):
         if not self.url:
             raise TypeError("Reference must have a url")
 
-    def normalized(self):
-        severities = sorted(self.severities)
-        return Reference(
-            reference_id=self.reference_id,
-            url=self.url,
-            severities=severities,
-            reference_type=self.reference_type,
-        )
+    def __lt__(self, other):
+        if not isinstance(other, Reference):
+            return NotImplemented
+        return self._cmp_key() < other._cmp_key()
+
+    # TODO: Add cache
+    def _cmp_key(self):
+        return (self.reference_id, self.reference_type, self.url, tuple(self.severities))
 
     def to_dict(self):
+        """Return a normalized dictionary representation"""
         return {
             "reference_id": self.reference_id,
             "reference_type": self.reference_type,
             "url": self.url,
-            "severities": [severity.to_dict() for severity in self.severities],
+            "severities": [severity.to_dict() for severity in sorted(self.severities)],
         }
 
     @classmethod
     def from_dict(cls, ref: dict):
         return cls(
-            reference_id=ref["reference_id"],
+            reference_id=str(ref["reference_id"]),
             reference_type=ref.get("reference_type") or "",
             url=ref["url"],
             severities=[
@@ -140,7 +155,8 @@ class NoAffectedPackages(Exception):
     """
 
 
-@dataclasses.dataclass(order=True, frozen=True)
+@functools.total_ordering
+@dataclasses.dataclass(eq=True)
 class AffectedPackage:
     """
     Relate a Package URL with a range of affected versions and a fixed version.
@@ -170,6 +186,19 @@ def get_fixed_purl(self):
             raise ValueError(f"Affected Package {self.package!r} does not have a fixed version")
         return update_purl_version(purl=self.package, version=str(self.fixed_version))
 
+    def __lt__(self, other):
+        if not isinstance(other, AffectedPackage):
+            return NotImplemented
+        return self._cmp_key() < other._cmp_key()
+
+    # TODO: Add cache
+    def _cmp_key(self):
+        return (
+            str(self.package),
+            str(self.affected_version_range or ""),
+            str(self.fixed_version or ""),
+        )
+
     @classmethod
     def merge(
         cls, affected_packages: Iterable

vulnerabilities/improvers/__init__.py

@@ -18,6 +18,7 @@
 from vulnerabilities.pipelines import enhance_with_kev
 from vulnerabilities.pipelines import enhance_with_metasploit
 from vulnerabilities.pipelines import flag_ghost_packages
+from vulnerabilities.pipelines import remove_duplicate_advisories
 
 IMPROVERS_REGISTRY = [
     valid_versions.GitHubBasicImprover,
@@ -45,6 +46,7 @@
     compute_package_version_rank.ComputeVersionRankPipeline,
     collect_commits.CollectFixCommitsPipeline,
     add_cvss31_to_CVEs.CVEAdvisoryMappingPipeline,
+    remove_duplicate_advisories.RemoveDuplicateAdvisoriesPipeline,
 ]
 
 IMPROVERS_REGISTRY = {

vulnerabilities/migrations/0089_alter_advisory_unique_content_id.py

@@ -0,0 +1,22 @@
+# Generated by Django 4.2.17 on 2025-02-27 07:47
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("vulnerabilities", "0088_fix_alpine_purl_type"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="advisory",
+            name="unique_content_id",
+            field=models.CharField(
+                blank=True,
+                help_text="A 64 character unique identifier for the content of the advisory since we use sha256 as hex",
+                max_length=64,
+            ),
+        ),
+    ]

vulnerabilities/models.py

@@ -53,6 +53,7 @@
 from vulnerabilities import utils
 from vulnerabilities.severity_systems import EPSS
 from vulnerabilities.severity_systems import SCORING_SYSTEMS
+from vulnerabilities.utils import compute_content_id
 from vulnerabilities.utils import normalize_purl
 from vulnerabilities.utils import purl_to_dict
 from vulnerablecode import __version__ as VULNERABLECODE_VERSION
@@ -1315,8 +1316,9 @@ class Advisory(models.Model):
     """
 
     unique_content_id = models.CharField(
-        max_length=32,
+        max_length=64,
         blank=True,
+        help_text="A 64 character unique identifier for the content of the advisory since we use sha256 as hex",
     )
     aliases = models.JSONField(blank=True, default=list, help_text="A list of alias strings")
     summary = models.TextField(
@@ -1357,16 +1359,8 @@ class Meta:
         ordering = ["aliases", "date_published", "unique_content_id"]
 
     def save(self, *args, **kwargs):
-        checksum = hashlib.md5()
-        for field in (
-            self.summary,
-            self.affected_packages,
-            self.references,
-            self.weaknesses,
-        ):
-            value = json.dumps(field, separators=(",", ":")).encode("utf-8")
-            checksum.update(value)
-        self.unique_content_id = checksum.hexdigest()
+        advisory_data = self.to_advisory_data()
+        self.unique_content_id = compute_content_id(advisory_data, include_metadata=False)
         super().save(*args, **kwargs)
 
     def to_advisory_data(self) -> "AdvisoryData":

vulnerabilities/pipelines/remove_duplicate_advisories.py

@@ -0,0 +1,110 @@
+#
+# 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/aboutcode-org/vulnerablecode for support or download.
+# See https://aboutcode.org for more information about nexB OSS projects.
+#
+
+from aboutcode.pipeline import LoopProgress
+
+from vulnerabilities.models import Advisory
+from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.utils import compute_content_id
+
+
+class RemoveDuplicateAdvisoriesPipeline(VulnerableCodePipeline):
+    """Pipeline to compute new advisory content id and remove duplicate advisories based on their content."""
+
+    pipeline_id = "remove_duplicate_advisories"
+
+    @classmethod
+    def steps(cls):
+        return (cls.remove_duplicates,)
+
+    def remove_duplicates(self):
+        """
+        Recompute the content ID and remove duplicate advisories, keeping the oldest one.
+        """
+
+        advisories_count = Advisory.objects.all().count()
+        self.log(f"Computing new content id for {advisories_count} and removing duplicates.")
+
+        update_batch_size = 500
+        delete_batch_size = 5000
+        chunk_size = 5000
+        deleted_advisories_count = 0
+        updated_advisories_count = 0
+        duplicate_advisory_ids = []
+        advisories_to_update = []
+        content_ids = set()
+
+        advisories = Advisory.objects.all().order_by("id").iterator(chunk_size=chunk_size)
+
+        progress = LoopProgress(
+            total_iterations=advisories_count,
+            logger=self.log,
+            progress_step=1,
+        )
+
+        for advisory in progress.iter(advisories):
+            content_id = compute_content_id(advisory.to_advisory_data())
+
+            if content_id in content_ids:
+                duplicate_advisory_ids.append(advisory.id)
+            else:
+                content_ids.add(content_id)
+                if advisory.unique_content_id != content_id:
+                    advisory.unique_content_id = content_id
+                    advisories_to_update.append(advisory)
+
+            if len(duplicate_advisory_ids) > delete_batch_size:
+                deleted_advisories_count += delete_advisories(
+                    advisory_ids=duplicate_advisory_ids,
+                    logger=self.log,
+                )
+                duplicate_advisory_ids.clear()
+
+            if len(advisories_to_update) > update_batch_size:
+                updated_advisories_count += bulk_update_advisories(
+                    advisories=advisories_to_update,
+                    fields=["unique_content_id"],
+                    logger=self.log,
+                )
+                advisories_to_update.clear()
+
+        deleted_advisories_count += delete_advisories(
+            advisory_ids=duplicate_advisory_ids,
+            logger=self.log,
+        )
+        updated_advisories_count += bulk_update_advisories(
+            advisories=advisories_to_update,
+            fields=["unique_content_id"],
+            logger=self.log,
+        )
+
+        self.log(f"Removed {deleted_advisories_count} duplicates advisories.")
+        self.log(f"Updated content id for {deleted_advisories_count} advisories.")
+
+
+def bulk_update_advisories(advisories, fields, logger):
+    item_count = 0
+    if advisories:
+        try:
+            Advisory.objects.bulk_update(objs=advisories, fields=fields)
+            item_count += len(advisories)
+        except Exception as e:
+            logger(f"Error updating Advisory: {e}")
+    return item_count
+
+
+def delete_advisories(advisory_ids, logger):
+    item_count = 0
+    if advisory_ids:
+        try:
+            Advisory.objects.filter(id__in=advisory_ids).delete()
+            item_count += len(advisory_ids)
+        except Exception as e:
+            logger(f"Error deleting Advisory: {e}")
+    return item_count

vulnerabilities/severity_systems.py

@@ -42,6 +42,9 @@ def compute(self, scoring_elements: str) -> str:
     def get(self, scoring_elements: str):
         return NotImplementedError
 
+    def __str__(self):
+        return f"{self.identifier}"
+
 
 @dataclasses.dataclass(order=True)
 class Cvssv2ScoringSystem(ScoringSystem):