Optimize creation of TODO and advisory TODO M2M relations

keshav-space · keshav-space · commit 550f72585b8b · 2025-06-27T22:51:17.000+05:30
Signed-off-by: Keshav Priyadarshi &lt;git@keshav.space&gt;
diff --git a/vulnerabilities/migrations/0093_advisorytodo_todorelatedadvisory_and_more.py b/vulnerabilities/migrations/0093_advisorytodo_todorelatedadvisory_and_more.py
@@ -1,6 +1,7 @@
-# Generated by Django 4.2.20 on 2025-06-03 18:13
+# Generated by Django 4.2.22 on 2025-06-27 15:59
 
 from django.db import migrations, models
+import django.db.models.deletion
 
 
 class Migration(migrations.Migration):
@@ -90,17 +91,47 @@ class Migration(migrations.Migration):
                         blank=True, help_text="Additional detail on how this TODO was resolved."
                     ),
                 ),
+            ],
+        ),
+        migrations.CreateModel(
+            name="ToDoRelatedAdvisory",
+            fields=[
                 (
-                    "advisories",
-                    models.ManyToManyField(
-                        help_text="Advisory/ies where this TODO is applicable.",
-                        related_name="advisory_todos",
-                        to="vulnerabilities.advisory",
+                    "id",
+                    models.AutoField(
+                        auto_created=True, primary_key=True, serialize=False, verbose_name="ID"
+                    ),
+                ),
+                (
+                    "advisory",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE, to="vulnerabilities.advisory"
+                    ),
+                ),
+                (
+                    "todo",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to="vulnerabilities.advisorytodo",
                     ),
                 ),
             ],
             options={
-                "unique_together": {("related_advisories_id", "issue_type")},
+                "unique_together": {("todo", "advisory")},
             },
         ),
+        migrations.AddField(
+            model_name="advisorytodo",
+            name="advisories",
+            field=models.ManyToManyField(
+                help_text="Advisory/ies where this TODO is applicable.",
+                related_name="advisory_todos",
+                through="vulnerabilities.ToDoRelatedAdvisory",
+                to="vulnerabilities.advisory",
+            ),
+        ),
+        migrations.AlterUniqueTogether(
+            name="advisorytodo",
+            unique_together={("related_advisories_id", "issue_type")},
+        ),
     ]
diff --git a/vulnerabilities/models.py b/vulnerabilities/models.py
@@ -2286,16 +2286,19 @@ class AdvisoryToDo(models.Model):
     # to avoid creating duplicate issue for same set of advisories,
     related_advisories_id = models.CharField(
         max_length=40,
-        blank=False,
-        null=False,
         help_text="SHA1 digest of the unique_content_id field of the applicable advisories.",
     )
 
+    advisories = models.ManyToManyField(
+        Advisory,
+        through="ToDoRelatedAdvisory",
+        related_name="advisory_todos",
+        help_text="Advisory/ies where this TODO is applicable.",
+    )
+
     issue_type = models.CharField(
         max_length=50,
         choices=ISSUE_TYPE_CHOICES,
-        blank=False,
-        null=False,
         db_index=True,
         help_text="Select the issue that needs to be addressed from the available options.",
     )
@@ -2305,12 +2308,6 @@ class AdvisoryToDo(models.Model):
         help_text="Additional details about the issue.",
     )
 
-    advisories = models.ManyToManyField(
-        Advisory,
-        related_name="advisory_todos",
-        help_text="Advisory/ies where this TODO is applicable.",
-    )
-
     created_at = models.DateTimeField(
         auto_now_add=True,
         help_text="Timestamp indicating when this TODO was created.",
@@ -2339,3 +2336,18 @@ class Meta:
     def save(self, *args, **kwargs):
         self.full_clean()
         return super().save(*args, **kwargs)
+
+
+class ToDoRelatedAdvisory(models.Model):
+    todo = models.ForeignKey(
+        AdvisoryToDo,
+        on_delete=models.CASCADE,
+    )
+
+    advisory = models.ForeignKey(
+        Advisory,
+        on_delete=models.CASCADE,
+    )
+
+    class Meta:
+        unique_together = ("todo", "advisory")
diff --git a/vulnerabilities/pipelines/compute_advisory_todo.py b/vulnerabilities/pipelines/compute_advisory_todo.py
@@ -11,10 +11,12 @@
 import json
 
 from aboutcode.pipeline import LoopProgress
+from django.utils import timezone
 
 from vulnerabilities.models import Advisory
 from vulnerabilities.models import AdvisoryToDo
 from vulnerabilities.models import Alias
+from vulnerabilities.models import ToDoRelatedAdvisory
 from vulnerabilities.pipelines import VulnerableCodePipeline
 from vulnerabilities.pipes.advisory import advisories_checksum
 
@@ -32,8 +34,14 @@ def steps(cls):
         )
 
     def compute_individual_advisory_todo(self):
-        advisories = Advisory.objects.all().iterator(chunk_size=5000)
-        advisories_count = Advisory.objects.all().count()
+        """Create ToDos for missing summary, affected and fixed packages."""
+
+        advisories = Advisory.objects.all()
+        advisories_count = advisories.count()
+        advisory_relation_to_create = {}
+        todo_to_create = []
+        new_todos_count = 0
+        batch_size = 5000
 
         self.log(
             f"Checking missing summary, affected and fixed packages in {advisories_count} Advisories"
@@ -43,23 +51,48 @@ def compute_individual_advisory_todo(self):
             logger=self.log,
             progress_step=1,
         )
-        for advisory in progress.iter(advisories):
+        for advisory in progress.iter(advisories.iterator(chunk_size=5000)):
             advisory_todo_id = advisories_checksum(advisories=advisory)
             check_missing_summary(
                 advisory=advisory,
                 todo_id=advisory_todo_id,
-                logger=self.log,
+                todo_to_create=todo_to_create,
+                advisory_relation_to_create=advisory_relation_to_create,
             )
 
             check_missing_affected_and_fixed_by_packages(
                 advisory=advisory,
                 todo_id=advisory_todo_id,
-                logger=self.log,
+                todo_to_create=todo_to_create,
+                advisory_relation_to_create=advisory_relation_to_create,
             )
 
+            if len(todo_to_create) > batch_size:
+                new_todos_count += bulk_create_with_m2m(
+                    todos=todo_to_create,
+                    advisories=advisory_relation_to_create,
+                    logger=self.log,
+                )
+                advisory_relation_to_create.clear()
+                todo_to_create.clear()
+
+        new_todos_count += bulk_create_with_m2m(
+            todos=todo_to_create,
+            advisories=advisory_relation_to_create,
+            logger=self.log,
+        )
+
     def detect_conflicting_advisories(self):
+        """
+        Create ToDos for advisories with conflicting opinions on fixed and affected
+        package versions for a vulnerability.
+        """
         aliases = Alias.objects.filter(alias__istartswith="cve")
         aliases_count = aliases.count()
+        advisory_relation_to_create = {}
+        todo_to_create = []
+        new_todos_count = 0
+        batch_size = 5000
 
         self.log(f"Cross validating advisory affected and fixed package for {aliases_count} CVEs")
 
@@ -73,24 +106,50 @@ def detect_conflicting_advisories(self):
                 advisory_todos__issue_type="MISSING_AFFECTED_AND_FIXED_BY_PACKAGES"
             ).distinct()
 
-            check_conflicting_affected_and_fixed_by_packages(
+            check_conflicting_affected_and_fixed_by_packages_for_alias(
                 advisories=advisories,
                 cve=alias,
-                logger=self.log,
+                todo_to_create=todo_to_create,
+                advisory_relation_to_create=advisory_relation_to_create,
             )
 
+            if len(todo_to_create) > batch_size:
+                new_todos_count += bulk_create_with_m2m(
+                    todos=todo_to_create,
+                    advisories=advisory_relation_to_create,
+                    logger=self.log,
+                )
+                advisory_relation_to_create.clear()
+                todo_to_create.clear()
+
+        new_todos_count += bulk_create_with_m2m(
+            todos=todo_to_create,
+            advisories=advisory_relation_to_create,
+            logger=self.log,
+        )
 
-def check_missing_summary(advisory, todo_id, logger=None):
+
+def check_missing_summary(
+    advisory,
+    todo_id,
+    todo_to_create,
+    advisory_relation_to_create,
+):
     if not advisory.summary:
-        todo, created = AdvisoryToDo.objects.get_or_create(
+        todo = AdvisoryToDo(
             related_advisories_id=todo_id,
             issue_type="MISSING_SUMMARY",
         )
-        if created:
-            todo.advisories.add(advisory)
+        advisory_relation_to_create[todo_id] = [advisory]
+        todo_to_create.append(todo)
 
 
-def check_missing_affected_and_fixed_by_packages(advisory, todo_id, logger=None):
+def check_missing_affected_and_fixed_by_packages(
+    advisory,
+    todo_id,
+    todo_to_create,
+    advisory_relation_to_create,
+):
     """
     Check for missing affected or fixed-by packages in the advisory
     and create appropriate AdvisoryToDo.
@@ -121,15 +180,21 @@ def check_missing_affected_and_fixed_by_packages(advisory, todo_id, logger=None)
         issue_type = "MISSING_AFFECTED_PACKAGE"
     elif not has_fixed_package:
         issue_type = "MISSING_FIXED_BY_PACKAGE"
-    todo, created = AdvisoryToDo.objects.get_or_create(
+
+    todo = AdvisoryToDo(
         related_advisories_id=todo_id,
         issue_type=issue_type,
     )
-    if created:
-        todo.advisories.add(advisory)
+    todo_to_create.append(todo)
+    advisory_relation_to_create[todo_id] = [advisory]
 
 
-def check_conflicting_affected_and_fixed_by_packages(advisories, cve, logger=None):
+def check_conflicting_affected_and_fixed_by_packages_for_alias(
+    advisories,
+    cve,
+    todo_to_create,
+    advisory_relation_to_create,
+):
     """
     Add appropriate AdvisoryToDo for conflicting affected/fixed packages.
 
@@ -222,15 +287,13 @@ def check_conflicting_affected_and_fixed_by_packages(advisories, cve, logger=Non
     messages.append("Comparison matrix:")
     messages.append(json.dumps(matrix, indent=2, default=list))
     todo_id = advisories_checksum(advisories)
-    todo, created = AdvisoryToDo.objects.get_or_create(
+    todo = AdvisoryToDo(
         related_advisories_id=todo_id,
         issue_type=issue_type,
-        defaults={
-            "issue_detail": "\n".join(messages),
-        },
+        issue_detail="\n".join(messages),
     )
-    if created:
-        todo.advisories.add(*advisories)
+    todo_to_create.append(todo)
+    advisory_relation_to_create[todo_id] = list(advisories)
 
 
 def initialize_sub_matrix(matrix, affected_purl, advisory):
@@ -245,3 +308,30 @@ def initialize_sub_matrix(matrix, affected_purl, advisory):
             matrix[affected_purl]["affected"][advisory_id] = set()
         if advisory not in matrix[affected_purl]["fixed"]:
             matrix[affected_purl]["fixed"][advisory_id] = set()
+
+
+def bulk_create_with_m2m(todos, advisories, logger):
+    """Bulk create ToDos and also bulk create M2M ToDo Advisory relationships."""
+    if not todos:
+        return 0
+
+    start_time = timezone.now()
+    try:
+        AdvisoryToDo.objects.bulk_create(objs=todos, ignore_conflicts=True)
+    except Exception as e:
+        logger(f"Error creating AdvisoryToDo: {e}")
+
+    new_todos = AdvisoryToDo.objects.filter(created_at__gte=start_time)
+
+    relations = [
+        ToDoRelatedAdvisory(todo=todo, advisory=advisory)
+        for todo in new_todos
+        for advisory in advisories[todo.related_advisories_id]
+    ]
+
+    try:
+        ToDoRelatedAdvisory.objects.bulk_create(relations)
+    except Exception as e:
+        logger(f"Error creating Advisory ToDo relations: {e}")
+
+    return new_todos.count()
diff --git a/vulnerabilities/tests/pipelines/test_compute_advisory_todo.py b/vulnerabilities/tests/pipelines/test_compute_advisory_todo.py
@@ -88,9 +88,10 @@ def test_advisory_todo_missing_summary(self):
         pipeline = ComputeToDo()
         pipeline.execute()
 
-        todos = AdvisoryToDo.objects.first()
+        todo = AdvisoryToDo.objects.first()
         self.assertEqual(1, AdvisoryToDo.objects.count())
-        self.assertEqual("MISSING_SUMMARY", todos.issue_type)
+        self.assertEqual("MISSING_SUMMARY", todo.issue_type)
+        self.assertEqual(1, todo.advisories.count())
 
     def test_advisory_todo_missing_fixed(self):
         date = datetime.now()
@@ -107,9 +108,10 @@ def test_advisory_todo_missing_fixed(self):
         pipeline = ComputeToDo()
         pipeline.execute()
 
-        todos = AdvisoryToDo.objects.first()
+        todo = AdvisoryToDo.objects.first()
         self.assertEqual(1, AdvisoryToDo.objects.count())
-        self.assertEqual("MISSING_FIXED_BY_PACKAGE", todos.issue_type)
+        self.assertEqual("MISSING_FIXED_BY_PACKAGE", todo.issue_type)
+        self.assertEqual(1, todo.advisories.count())
 
     def test_advisory_todo_missing_affected(self):
         date = datetime.now()
@@ -126,9 +128,10 @@ def test_advisory_todo_missing_affected(self):
         pipeline = ComputeToDo()
         pipeline.execute()
 
-        todos = AdvisoryToDo.objects.first()
+        todo = AdvisoryToDo.objects.first()
         self.assertEqual(1, AdvisoryToDo.objects.count())
-        self.assertEqual("MISSING_AFFECTED_PACKAGE", todos.issue_type)
+        self.assertEqual("MISSING_AFFECTED_PACKAGE", todo.issue_type)
+        self.assertEqual(1, todo.advisories.count())
 
     def test_advisory_todo_conflicting_fixed_affected(self):
         alias = Alias.objects.create(alias="CVE-0000-0000")
@@ -156,12 +159,15 @@ def test_advisory_todo_conflicting_fixed_affected(self):
         )
         adv2.aliases.add(alias)
 
+        self.assertEqual(0, AdvisoryToDo.objects.count())
         pipeline = ComputeToDo()
         pipeline.execute()
 
-        todos = AdvisoryToDo.objects.first()
+        todo = AdvisoryToDo.objects.first()
         self.assertEqual(1, AdvisoryToDo.objects.count())
-        self.assertEqual("CONFLICTING_AFFECTED_AND_FIXED_BY_PACKAGES", todos.issue_type)
+        self.assertEqual("CONFLICTING_AFFECTED_AND_FIXED_BY_PACKAGES", todo.issue_type)
         self.assertIn(
-            "CVE-0000-0000: pkg:npm/package1 with conflicting fixed version", todos.issue_detail
+            "CVE-0000-0000: pkg:npm/package1 with conflicting fixed version", todo.issue_detail
         )
+        self.assertEqual(2, todo.advisories.count())
+        self.assertEqual(todo, adv2.advisory_todos.first())
