Make the migration reversible

arkindex/training/migrations/0007_datasetset_model.py

@@ -27,42 +27,68 @@ class Migration(migrations.Migration):
             model_name="datasetset",
             constraint=models.UniqueConstraint(fields=("dataset", "name"), name="unique_dataset_sets"),
         ),
+        # Make the old set name and dataset ID fields nullable
+        # so that they can be filled in when rolling the migration back
+        migrations.AlterField(
+            model_name="datasetelement",
+            name="set",
+            field=models.CharField(max_length=50, null=True),
+        ),
+        migrations.AlterField(
+            model_name="datasetelement",
+            name="dataset",
+            field=models.ForeignKey(on_delete=django.db.models.deletion.DO_NOTHING, related_name="elements", to="training.dataset", null=True),
+        ),
         migrations.AddField(
             model_name="datasetelement",
+            # Temporary name, because the `set` column already existed as the set name.
+            # This is referred to as `set_id_id` in the SQL migration,
+            # and renamed to `set` afterwards.
             name="set_id",
             field=models.ForeignKey(null=True, on_delete=django.db.models.deletion.DO_NOTHING, related_name="dataset_elements", to="training.datasetset"),
         ),
         migrations.RunSQL(
-            """
-            INSERT INTO training_datasetset (id, dataset_id, name)
-            SELECT uuid_generate_v4(), ds.id, ds.set
-            FROM (
-                SELECT DISTINCT id, unnest(sets) AS set
-                FROM training_dataset
-            ) ds
-            """,
-            reverse_sql=migrations.RunSQL.noop,
-        ),
-        migrations.RunSQL(
-            """
-            INSERT INTO training_datasetset (id, dataset_id, name)
-            SELECT uuid_generate_v4(), sets.dataset_id, sets.set
-            FROM (
-                SELECT DISTINCT dataset_id, set
-                FROM training_datasetelement
-            ) sets
-            ON CONFLICT DO NOTHING
-            """,
-            reverse_sql=migrations.RunSQL.noop,
-        ),
-        migrations.RunSQL(
-            """
-            UPDATE training_datasetelement de
-            SET set_id_id = ds.id
-            FROM training_datasetset ds
-            WHERE de.dataset_id = ds.dataset_id AND de.set = ds.name
-            """,
-            reverse_sql=migrations.RunSQL.noop,
+            [
+                """
+                INSERT INTO training_datasetset (id, dataset_id, name)
+                SELECT uuid_generate_v4(), ds.id, ds.set
+                FROM (
+                    SELECT DISTINCT id, unnest(sets) AS set
+                    FROM training_dataset
+                ) ds
+                """,
+                """
+                INSERT INTO training_datasetset (id, dataset_id, name)
+                SELECT uuid_generate_v4(), sets.dataset_id, sets.set
+                FROM (
+                    SELECT DISTINCT dataset_id, set
+                    FROM training_datasetelement
+                ) sets
+                ON CONFLICT DO NOTHING
+                """,
+                """
+                UPDATE training_datasetelement de
+                SET set_id_id = ds.id
+                FROM training_datasetset ds
+                WHERE de.dataset_id = ds.dataset_id AND de.set = ds.name
+                """,
+            ],
+            reverse_sql=[
+                """
+                UPDATE training_dataset
+                SET sets = ARRAY(
+                    SELECT name
+                    FROM training_datasetset
+                    WHERE dataset_id = training_dataset.id
+                )
+                """,
+                """
+                UPDATE training_datasetelement de
+                SET dataset_id = ds.dataset_id, set = ds.name
+                FROM training_datasetset ds
+                WHERE ds.id = de.set_id_id
+                """,
+            ],
         ),
         migrations.RemoveConstraint(
             model_name="datasetelement",