Add migration to populate WorkflowState's content_type and base_content_type fields

wagtail/migrations/0081_populate_workflowstate_content_type.py

@@ -0,0 +1,41 @@
+# Generated by Django 4.2.dev20221212151407 on 2022-12-12 15:16
+
+from django.db import migrations, models
+from django.db.models.functions import Cast
+
+
+def populate_workflowstate_content_type(apps, schema_editor):
+    ContentType = apps.get_model("contenttypes.ContentType")
+    WorkflowState = apps.get_model("wagtailcore.WorkflowState")
+    Page = apps.get_model("wagtailcore.Page")
+
+    page_type = ContentType.objects.get(app_label="wagtailcore", model="page")
+    content_type_id = models.Subquery(
+        Page.objects.filter(
+            pk=Cast(models.OuterRef("object_id"), models.PositiveIntegerField())
+        ).values("content_type_id")
+    )
+
+    WorkflowState.objects.all().update(
+        base_content_type=page_type,
+        content_type_id=content_type_id,
+    )
+
+
+def empty_workflowstate_content_type(apps, schema_editor):
+    WorkflowState = apps.get_model("wagtailcore.WorkflowState")
+    WorkflowState.objects.all().update(base_content_type=None, content_type=None)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("wagtailcore", "0080_generic_workflowstate"),
+    ]
+
+    operations = [
+        migrations.RunPython(
+            populate_workflowstate_content_type,
+            empty_workflowstate_content_type,
+        )
+    ]