Merge pull request #1630 from aboutcode-org/fix-migration

Avoid memory exhaustion during data migration

vulnerabilities/migrations/0071_auto_20241007_1044.py

@@ -1,14 +1,22 @@
 from django.db import migrations, models
 import django.db.models.deletion
-from django.core.validators import MaxValueValidator, MinValueValidator
-from vulnerabilities.improver import MAX_CONFIDENCE
+from aboutcode.pipeline import LoopProgress
 
 def split_packagerelatedvulnerability(apps, schema_editor):
     PackageRelatedVulnerability = apps.get_model('vulnerabilities', 'PackageRelatedVulnerability')
     FixingPackageRelatedVulnerability = apps.get_model('vulnerabilities', 'FixingPackageRelatedVulnerability')
     AffectedByPackageRelatedVulnerability = apps.get_model('vulnerabilities', 'AffectedByPackageRelatedVulnerability')
 
-    for prv in PackageRelatedVulnerability.objects.all():
+    obsolete_package_relation_query = PackageRelatedVulnerability.objects.all()
+    obsolete_package_relation_query_count = obsolete_package_relation_query.count()
+    print(f"\nMigrating {obsolete_package_relation_query_count:,d} old package vulnerability relationship.")
+
+    progress = LoopProgress(
+        total_iterations=obsolete_package_relation_query_count,
+        progress_step=1,
+        logger=print,
+        )
+    for prv in progress.iter(obsolete_package_relation_query.iterator(chunk_size=10000)):
         if prv.fix:
             FixingPackageRelatedVulnerability.objects.create(
                 package=prv.package,
@@ -29,7 +37,16 @@ def reverse_migration(apps, schema_editor):
     AffectedByPackageRelatedVulnerability = apps.get_model('vulnerabilities', 'AffectedByPackageRelatedVulnerability')
     PackageRelatedVulnerability = apps.get_model('vulnerabilities', 'PackageRelatedVulnerability')
 
-    for fpv in FixingPackageRelatedVulnerability.objects.all():
+    fixing_package_relation_query = FixingPackageRelatedVulnerability.objects.all()
+    fixing_package_relation_query_count = fixing_package_relation_query.count()
+    print(f"\nMigrating {fixing_package_relation_query_count:,d} FixingPackage to old relationship.")
+
+    progress = LoopProgress(
+        total_iterations=fixing_package_relation_query_count,
+        progress_step=1,
+        logger=print,
+        )
+    for fpv in progress.iter(fixing_package_relation_query.iterator(chunk_size=10000)):
         PackageRelatedVulnerability.objects.create(
             package=fpv.package,
             vulnerability=fpv.vulnerability,
@@ -38,7 +55,16 @@ def reverse_migration(apps, schema_editor):
             fix=True,
         )
 
-    for apv in AffectedByPackageRelatedVulnerability.objects.all():
+    affected_package_relation_query = AffectedByPackageRelatedVulnerability.objects.all()
+    affected_package_relation_query_count = affected_package_relation_query.count()
+    print(f"\nMigrating {affected_package_relation_query_count:,d} AffectedPackage to old relationship.")
+
+    progress = LoopProgress(
+        total_iterations=affected_package_relation_query_count,
+        progress_step=1,
+        logger=print,
+        )
+    for apv in progress.iter(affected_package_relation_query.iterator(chunk_size=10000)):
         PackageRelatedVulnerability.objects.create(
             package=apv.package,
             vulnerability=apv.vulnerability,