Data migration to backfill new category counts

temba/flows/migrations/0355_backfill_new_cat_counts.py

@@ -0,0 +1,50 @@
+# Generated by Django 5.1.4 on 2025-01-09 18:35
+
+import itertools
+
+from django.db import migrations, transaction
+from django.db.models import Sum
+
+
+def backfill_new_counts(apps, schema_editor):
+    Flow = apps.get_model("flows", "Flow")
+
+    flow_ids = list(Flow.objects.filter(is_active=True).order_by("id").values_list("id", flat=True))
+
+    print(f"Updating result counts for {len(flow_ids)} flows...")
+
+    num_backfilled = 0
+
+    for id_batch in itertools.batched(flow_ids, 500):
+        flows = Flow.objects.filter(id__in=id_batch).only("id").order_by("id")
+        for flow in flows:
+            backfill_for_flow(apps, flow)
+
+        num_backfilled += len(flows)
+        print(f"> updated counts for {num_backfilled} of {len(flow_ids)} flows")
+
+
+def backfill_for_flow(apps, flow) -> int:
+    FlowResultCount = apps.get_model("flows", "FlowResultCount")
+
+    to_create = []
+
+    counts = flow.category_counts.values("result_key", "category_name").annotate(total=Sum("count"))
+    for c in counts:
+        if c["category_name"] and c["total"] > 0:
+            to_create.append(
+                FlowResultCount(
+                    flow=flow, result=c["result_key"][:64], category=c["category_name"][:64], count=c["total"]
+                )
+            )
+
+    with transaction.atomic():
+        flow.result_counts.all().delete()
+        FlowResultCount.objects.bulk_create(to_create)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [("flows", "0354_flowresultcount")]
+
+    operations = [migrations.RunPython(backfill_new_counts, migrations.RunPython.noop)]

temba/flows/tests/test_migrations.py

@@ -0,0 +1,65 @@
+from temba.tests import MigrationTest
+
+
+class BackfillNewCategoryCountsTest(MigrationTest):
+    app = "flows"
+    migrate_from = "0354_flowresultcount"
+    migrate_to = "0355_backfill_new_cat_counts"
+
+    def setUpBeforeMigration(self, apps):
+        self.flow1 = self.create_flow("Flow 1")
+        self.flow2 = self.create_flow("Flow 2")
+
+        self.flow1.category_counts.create(
+            node_uuid="0043bf6f-f385-4ba3-80d7-4313771efa86",
+            result_key="color",
+            result_name="Color",
+            category_name="Red",
+            count=1,
+        )
+        self.flow1.category_counts.create(
+            node_uuid="2c230f54-a7f6-4a71-9f24-b3ba81734552",
+            result_key="color",
+            result_name="Color",
+            category_name="Red",
+            count=2,
+        )
+        self.flow1.category_counts.create(
+            node_uuid="0043bf6f-f385-4ba3-80d7-4313771efa86",
+            result_key="color",
+            result_name="Color",
+            category_name="Blue",
+            count=4,
+        )
+        self.flow1.category_counts.create(
+            node_uuid="f27d77e9-4ccb-4c36-8277-f5708822cf26",
+            result_key="name",
+            result_name="Name",
+            category_name="All Responses",
+            count=6,
+        )
+        self.flow2.category_counts.create(
+            node_uuid="daf35e03-b82e-4c12-8ff7-60f7d990eb05",
+            result_key="thing",
+            result_name="Thing",
+            category_name="Looooonnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnggggggggggggnnggggggggggggggggg",
+            count=2,
+        )
+        self.flow2.category_counts.create(
+            node_uuid="daf35e03-b82e-4c12-8ff7-60f7d990eb05",
+            result_key="thing",
+            result_name="Thing",
+            category_name="Zero",
+            count=0,
+        )
+
+    def test_migration(self):
+        def assert_counts(flow, expected):
+            actual = {}
+            for count in flow.result_counts.all():
+                actual[f"{count.result}/{count.category}"] = count.count
+
+            self.assertEqual(actual, expected)
+
+        assert_counts(self.flow1, {"color/Red": 3, "color/Blue": 4, "name/All Responses": 6})
+        assert_counts(self.flow2, {"thing/Looooonnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnggggggggggggnngggggggggg": 2})