Optimize alias migration using AdvisoryRelatedAlias

- Use AdvisoryRelatedAlias to manage advisory alias relation

Signed-off-by: Keshav Priyadarshi <[email protected]>

Author: keshav-space

vulnerabilities/migrations/0090_migrate_advisory_aliases.py

@@ -7,16 +7,32 @@
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
 
+from timeit import default_timer as timer
+
+import django.db.models.deletion
 from aboutcode.pipeline import LoopProgress
+from aboutcode.pipeline import humanize_time
 from django.db import migrations
 from django.db import models
-import django.db.models.deletion
 
 """
-Model and data migration for converting the Advisory aliases
-JSON field to a concrete M2M Advisory Alias relationship.
+Model and data migration to convert Advisory.aliases
+JSON field to a concrete M2M Advisory-Alias relationship.
+
+To achieve this following steps are executed in chronological order.
+ - Create AdvisoryRelatedAlias model for Advisory-Alias M2M relationship.
+ - Make unique_content_id non-nullable and a required field.
+ - Make Alias.vulnerability field nullable, as vulnerability may not 
+   exist for a corresponding alias during initial data collection.
+ - Rename existing Advisory.aliases JSON field to old_aliases.
+ - Create a new Advisory.aliases M2M relation through AdvisoryRelatedAlias model.
+ - Run a data migration to populate new M2M Advisory.aliases relation using 
+    Advisory.old_aliases data.
+ - Delete Advisory.old_aliases field.
+
 """
 
+
 def bulk_update(model, items, fields, logger):
     item_count = 0
     if items:
@@ -25,7 +41,17 @@ def bulk_update(model, items, fields, logger):
             item_count += len(items)
         except Exception as e:
             logger(f"Error updating Advisory: {e}")
-        items.clear()
+    return item_count
+
+
+def bulk_create(model, items, logger):
+    item_count = 0
+    if items:
+        try:
+            model.objects.bulk_create(objs=items)
+            item_count += len(items)
+        except Exception as e:
+            logger(f"Error creating AdvisoryRelatedAlias: {e}")
     return item_count
 
 
@@ -36,56 +62,122 @@ class Migration(migrations.Migration):
     ]
 
     def populate_new_advisory_aliases_field(apps, schema_editor):
+        """Populate the new Advisory.aliases relation using old_aliases JSON data."""
+        migration_start_time = timer()
         Advisory = apps.get_model("vulnerabilities", "Advisory")
         Alias = apps.get_model("vulnerabilities", "Alias")
+        AdvisoryRelatedAlias = apps.get_model("vulnerabilities", "AdvisoryRelatedAlias")
         advisories = Advisory.objects.all()
+        aliases = {i.alias: i for i in Alias.objects.all()}
 
-        chunk_size = 10000
+        chunk_size = 5000
         advisories_count = advisories.count()
-        print(f"\nPopulate new advisory aliases relationship.")
+        batch_size = 5000
+        relation_to_create = []
+        advisory_alias_relation_count = 0
         progress = LoopProgress(
             total_iterations=advisories_count,
             logger=print,
             progress_step=1,
         )
+        print(f"\nPopulate new advisory aliases relationship.")
         for advisory in progress.iter(advisories.iterator(chunk_size=chunk_size)):
-            aliases = Alias.objects.filter(alias__in=advisory.old_aliases)
-            advisory.aliases.set(aliases)
+            advisory_alias_relations = [
+                AdvisoryRelatedAlias(advisory=advisory, alias=aliases[alias])
+                for alias in advisory.old_aliases
+                if alias in aliases
+            ]
+            relation_to_create.extend(advisory_alias_relations)
+
+            if len(relation_to_create) > batch_size:
+                advisory_alias_relation_count += bulk_create(
+                    model=AdvisoryRelatedAlias,
+                    items=relation_to_create,
+                    logger=print,
+                )
+                relation_to_create.clear()
+
+        advisory_alias_relation_count += bulk_create(
+            model=AdvisoryRelatedAlias,
+            items=relation_to_create,
+            logger=print,
+        )
+        migration_run_time = timer() - migration_start_time
+        print(
+            f"\nSuccessfully created {advisory_alias_relation_count} advisory-alias relationship."
+        )
+        print(f"\nData Migration: completed in {humanize_time(migration_run_time)}")
 
     def reverse_populate_new_advisory_aliases_field(apps, schema_editor):
+        """Use the Advisory.aliases relation to populate old_aliases JSON field."""
+        migration_start_time = timer()
         Advisory = apps.get_model("vulnerabilities", "Advisory")
-        advisories = Advisory.objects.all()
+        advisories = Advisory.objects.prefetch_related("aliases").all()
 
         updated_advisory_count = 0
-        batch_size = 10000
-        chunk_size = 10000
-        updated_advisory = []
+        batch_size = 5000
+        chunk_size = 5000
+        advisory_to_update = []
         progress = LoopProgress(
             total_iterations=advisories.count(),
             logger=print,
             progress_step=1,
         )
+        print(f"\nReverse alias migration to M2M relation.")
         for advisory in progress.iter(advisories.iterator(chunk_size=chunk_size)):
             aliases = advisory.aliases.all()
             advisory.old_aliases = [alias.alias for alias in aliases]
-            updated_advisory.append(advisory)
+            advisory_to_update.append(advisory)
 
-            if len(updated_advisory) > batch_size:
+            if len(advisory_to_update) > batch_size:
                 updated_advisory_count += bulk_update(
                     model=Advisory,
-                    items=updated_advisory,
+                    items=advisory_to_update,
                     fields=["old_aliases"],
                     logger=print,
                 )
+                advisory_to_update.clear()
 
         updated_advisory_count += bulk_update(
             model=Advisory,
-            items=updated_advisory,
+            items=advisory_to_update,
             fields=["old_aliases"],
             logger=print,
         )
 
+        migration_run_time = timer() - migration_start_time
+        print(
+            f"\nSuccessfully reversed the alias relationship for {updated_advisory_count} advisories."
+        )
+        print(f"\nData Migration: completed in {humanize_time(migration_run_time)}")
+
     operations = [
+        migrations.CreateModel(
+            name="AdvisoryRelatedAlias",
+            fields=[
+                (
+                    "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"
+                    ),
+                ),
+                (
+                    "alias",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE, to="vulnerabilities.alias"
+                    ),
+                ),
+            ],
+            options={
+                "unique_together": {("advisory", "alias")},
+            },
+        ),
         migrations.AlterField(
             model_name="advisory",
             name="unique_content_id",
@@ -96,8 +188,6 @@ def reverse_populate_new_advisory_aliases_field(apps, schema_editor):
                 null=False,
             ),
         ),
-
-        # Make vulnerability relation optional
         migrations.AlterField(
             model_name="alias",
             name="vulnerability",
@@ -109,8 +199,6 @@ def reverse_populate_new_advisory_aliases_field(apps, schema_editor):
                 to="vulnerabilities.vulnerability",
             ),
         ),
-
-        # Rename aliases field to old_aliases
         migrations.AlterModelOptions(
             name="advisory",
             options={"ordering": ["date_published", "unique_content_id"]},
@@ -127,14 +215,16 @@ def reverse_populate_new_advisory_aliases_field(apps, schema_editor):
         migrations.AddField(
             model_name="advisory",
             name="aliases",
-            field=models.ManyToManyField(related_name="advisories", to="vulnerabilities.alias"),
+            field=models.ManyToManyField(
+                related_name="advisories",
+                through="vulnerabilities.AdvisoryRelatedAlias",
+                to="vulnerabilities.alias",
+            ),
         ),
-        # Populate the new M2M aliases relation
         migrations.RunPython(
             code=populate_new_advisory_aliases_field,
             reverse_code=reverse_populate_new_advisory_aliases_field,
         ),
-        # Delete JSON aliases field
         migrations.RemoveField(
             model_name="advisory",
             name="old_aliases",

vulnerabilities/models.py

@@ -1325,6 +1325,7 @@ class Advisory(models.Model):
     )
     aliases = models.ManyToManyField(
         Alias,
+        through="AdvisoryRelatedAlias",
         related_name="advisories",
     )
     summary = models.TextField(
@@ -1386,6 +1387,21 @@ def to_advisory_data(self) -> "AdvisoryData":
         )
 
 
+class AdvisoryRelatedAlias(models.Model):
+    advisory = models.ForeignKey(
+        Advisory,
+        on_delete=models.CASCADE,
+    )
+
+    alias = models.ForeignKey(
+        Alias,
+        on_delete=models.CASCADE,
+    )
+
+    class Meta:
+        unique_together = ("advisory", "alias")
+
+
 UserModel = get_user_model()