Fix populate_diff_id_fields management command

django/cantusdb_project/main_app/management/commands/populate_diff_id_fields.py

@@ -2,44 +2,42 @@
 from django.core.management.base import BaseCommand
 from django.db.models import Q
 from typing import Optional
+from django.db import transaction
 
 
 class Command(BaseCommand):
+    @transaction.atomic
     def handle(self, *args, **kwargs):
-        CHUNK_SIZE = 500
         chants = Chant.objects.filter(
             Q(differentiae_database__isnull=False) & Q(diff_db__isnull=True)
-        )
-        chants_count = chants.count()
-        start_index = 0
+        ).iterator(chunk_size=500)
+        chants_total = Chant.objects.filter(
+            Q(differentiae_database__isnull=False) & Q(diff_db__isnull=True)
+        ).count()
         count = 0
-        while start_index <= chants_count:
-            self.stdout.write(f"processing chunk with {start_index=}")
-            chunk = chants[start_index : start_index + CHUNK_SIZE]
 
-            for chant in chunk:
-                try:
-                    differentia_id: Optional[str] = chant.differentiae_database
-                    differentia = Differentia.objects.get(differentia_id=differentia_id)
-                    if differentia:
-                        chant.diff_db = differentia
-                    else:
-                        # If the Differentia doesn't exist, create a new one
-                        differentia = Differentia(
-                            differentia_id=differentia_id,
-                        )
-                        differentia.save()
-                        chant.diff_db = differentia
-                    chant.save()
-                except Differentia.DoesNotExist:
-                    print(f"Differentia not found for chant: {chant}")
-                count += 1
-                if count % 100 == 0:
-                    print(
-                        f"------------------ {count} of {chants_count} chants updated ------------------"
+        for chant in chants:
+            try:
+                differentia_id: Optional[str] = chant.differentiae_database
+                differentia = Differentia.objects.get(differentia_id=differentia_id)
+            except Differentia.DoesNotExist:
+                # If the Differentia doesn't exist, create a new one
+                differentia = Differentia(differentia_id=differentia_id)
+                differentia.save()
+                self.stdout.write(
+                    self.style.WARNING(f"Differentia created for chant: {chant}")
+                )
+
+            chant.diff_db = differentia
+            chant.save()
+
+            count += 1
+            if count % 100 == 0:
+                self.stdout.write(
+                    self.style.SUCCESS(
+                        f"------------------ {count} of {chants_total} chants updated ------------------"
                     )
-            del chunk  # make sure we don't use too much RAM
-            start_index += CHUNK_SIZE
+                )
 
         self.stdout.write(
             self.style.SUCCESS("Success! Command has run to completion.\n")