Migrate Advisory aliases field to M2M relationship

vulnerabilities/import_runner.py

@@ -15,6 +15,7 @@
 
 from django.core.exceptions import ValidationError
 from django.db import transaction
+from django.db.models.query import QuerySet
 
 from vulnerabilities.importer import AdvisoryData
 from vulnerabilities.importer import Importer
@@ -96,23 +97,29 @@ def process_advisories(
         Insert advisories into the database
         Return the number of inserted advisories.
         """
+        from vulnerabilities.pipes.advisory import get_or_create_aliases
+        from vulnerabilities.utils import compute_content_id
+
         count = 0
         advisories = []
         for data in advisory_datas:
+            content_id = compute_content_id(advisory_data=data)
             try:
+                aliases = get_or_create_aliases(aliases=data.aliases)
                 obj, created = Advisory.objects.get_or_create(
-                    aliases=data.aliases,
-                    summary=data.summary,
-                    affected_packages=[pkg.to_dict() for pkg in data.affected_packages],
-                    references=[ref.to_dict() for ref in data.references],
-                    date_published=data.date_published,
-                    weaknesses=data.weaknesses,
+                    unique_content_id=content_id,
+                    url=data.url,
                     defaults={
+                        "summary": data.summary,
+                        "affected_packages": [pkg.to_dict() for pkg in data.affected_packages],
+                        "references": [ref.to_dict() for ref in data.references],
+                        "date_published": data.date_published,
+                        "weaknesses": data.weaknesses,
                         "created_by": importer_name,
                         "date_collected": datetime.datetime.now(tz=datetime.timezone.utc),
                     },
-                    url=data.url,
                 )
+                obj.aliases.add(*aliases)
                 if not obj.date_imported:
                     advisories.append(obj)
             except Exception as e:
@@ -148,6 +155,8 @@ def process_inferences(inferences: List[Inference], advisory: Advisory, improver
     erroneous. Also, the atomic transaction for every advisory and its
     inferences makes sure that date_imported of advisory is consistent.
     """
+    from vulnerabilities.pipes.advisory import get_or_create_aliases
+
     inferences_processed_count = 0
 
     if not inferences:
@@ -157,9 +166,10 @@ def process_inferences(inferences: List[Inference], advisory: Advisory, improver
     logger.info(f"Improving advisory id: {advisory.id}")
 
     for inference in inferences:
+        aliases = get_or_create_aliases(inference.aliases)
         vulnerability = get_or_create_vulnerability_and_aliases(
             vulnerability_id=inference.vulnerability_id,
-            aliases=inference.aliases,
+            aliases=aliases,
             summary=inference.summary,
             advisory=advisory,
         )
@@ -265,14 +275,13 @@ def create_valid_vulnerability_reference(url, reference_id=None):
 
 
 def get_or_create_vulnerability_and_aliases(
-    aliases: List[str], vulnerability_id=None, summary=None, advisory=None
+    aliases: QuerySet, vulnerability_id=None, summary=None, advisory=None
 ):
     """
     Get or create vulnerabilitiy and aliases such that all existing and new
     aliases point to the same vulnerability
     """
-    aliases = set(alias.strip() for alias in aliases if alias and alias.strip())
-    new_alias_names, existing_vulns = get_vulns_for_aliases_and_get_new_aliases(aliases)
+    new_aliases, existing_vulns = get_vulns_for_aliases_and_get_new_aliases(aliases)
 
     # All aliases must point to the same vulnerability
     vulnerability = None
@@ -310,11 +319,11 @@ def get_or_create_vulnerability_and_aliases(
         #         f"Inconsistent summary for {vulnerability.vulnerability_id}. "
         #         f"Existing: {vulnerability.summary!r}, provided: {summary!r}"
         #     )
-        associate_vulnerability_with_aliases(vulnerability=vulnerability, aliases=new_alias_names)
+        associate_vulnerability_with_aliases(vulnerability=vulnerability, aliases=new_aliases)
     else:
         try:
             vulnerability = create_vulnerability_and_add_aliases(
-                aliases=new_alias_names, summary=summary
+                aliases=new_aliases, summary=summary
             )
             importer_name = get_importer_name(advisory)
             VulnerabilityChangeLog.log_import(
@@ -324,24 +333,22 @@ def get_or_create_vulnerability_and_aliases(
             )
         except Exception as e:
             logger.error(
-                f"Cannot create vulnerability with summary {summary!r} and {new_alias_names!r} {e!r}.\n{traceback_format_exc()}."
+                f"Cannot create vulnerability with summary {summary!r} and {new_aliases!r} {e!r}.\n{traceback_format_exc()}."
             )
             return
 
     return vulnerability
 
 
-def get_vulns_for_aliases_and_get_new_aliases(aliases):
+def get_vulns_for_aliases_and_get_new_aliases(aliases: QuerySet):
     """
     Return ``new_aliases`` that are not in the database and
     ``existing_vulns`` that point to the given ``aliases``.
     """
-    new_aliases = set(aliases)
-    existing_vulns = set()
-    for alias in Alias.objects.filter(alias__in=aliases):
-        existing_vulns.add(alias.vulnerability)
-        new_aliases.remove(alias.alias)
-    return new_aliases, existing_vulns
+    new_aliases = aliases.filter(vulnerability__isnull=True)
+    existing_vulns = [alias.vulnerability for alias in aliases.filter(vulnerability__isnull=False)]
+
+    return new_aliases, list(set(existing_vulns))
 
 
 @transaction.atomic
@@ -360,7 +367,5 @@ def create_vulnerability_and_add_aliases(aliases, summary):
 
 
 def associate_vulnerability_with_aliases(aliases, vulnerability):
-    for alias_name in aliases:
-        alias = Alias(alias=alias_name, vulnerability=vulnerability)
-        alias.save()
-        logger.info(f"New alias for {vulnerability!r}: {alias_name}")
+    aliases.update(vulnerability=vulnerability)
+    logger.info(f"New alias for {vulnerability!r}: {aliases}")

vulnerabilities/importer.py

@@ -103,6 +103,8 @@ class Reference:
     def __post_init__(self):
         if not self.url:
             raise TypeError("Reference must have a url")
+        if self.reference_id and not isinstance(self.reference_id, str):
+            self.reference_id = str(self.reference_id)
 
     def __lt__(self, other):
         if not isinstance(other, Reference):

vulnerabilities/improvers/vulnerability_status.py

@@ -10,7 +10,6 @@
 from typing import Iterable
 from urllib.parse import urljoin
 
-from django.db.models import Q
 from django.db.models.query import QuerySet
 
 from vulnerabilities.importer import AdvisoryData
@@ -37,10 +36,8 @@ class VulnerabilityStatusImprover(Improver):
 
     @property
     def interesting_advisories(self) -> QuerySet:
-        return (
-            Advisory.objects.filter(Q(created_by=NVDImporterPipeline.pipeline_id))
-            .distinct("aliases")
-            .paginated()
+        return Advisory.objects.filter(created_by=NVDImporterPipeline.pipeline_id).iterator(
+            chunk_size=5000
         )
 
     def get_inferences(self, advisory_data: AdvisoryData) -> Iterable[Inference]: