1 files changed, 51 insertions, 0 deletions
diff --git a/archaeological_finds/migrations/0137_data_migration_treatment_status_inputstatus.py b/archaeological_finds/migrations/0137_data_migration_treatment_status_inputstatus.py
new file mode 100644
index 000000000..b1f2077a3
--- /dev/null
+++ b/archaeological_finds/migrations/0137_data_migration_treatment_status_inputstatus.py
@@ -0,0 +1,51 @@
+import os
+
+from django.db import migrations
+from django.core.management import call_command
+
+
+def load_data(apps, __):
+    TreatmentStatus = apps.get_model("archaeological_finds", "treatmentstatus")
+    if not TreatmentStatus.objects.count():
+        migration = "0137_data_migration_treatmentstatus.json"
+        json_path = os.sep.join(os.path.abspath(__file__).split(os.sep)[:-1] + [migration])
+        call_command("loaddata", json_path)
+
+    TreatmentInputStatus = apps.get_model("archaeological_finds", "treatmentinputstatus")
+    if not TreatmentInputStatus.objects.count():
+        migration = "0137_data_migration_treatmentinputstatus.json"
+        json_path = os.sep.join(os.path.abspath(__file__).split(os.sep)[:-1] + [migration])
+        call_command("loaddata", json_path)
+
+    input_status_executed = TreatmentInputStatus.objects.filter(executed=True).all()[0]
+    input_status_draft = TreatmentInputStatus.objects.filter(executed=False).all()[0]
+
+    Treatment = apps.get_model("archaeological_finds", "treatment")
+    try:
+        TreatmentState = apps.get_model("archaeological_finds", "treatmentstate")
+    except LookupError:
+        # treatment state model deleted
+        return
+    for treatment_state in TreatmentState.objects.all():
+        q = TreatmentStatus.objects.filter(txt_idx=treatment_state.txt_idx)
+        if not q.exists():
+            print(f"Equivalence for treatment state: '{treatment_state.txt_idx}' not found")
+            continue
+        equiv = q.all()[0]
+        q = Treatment.objects.filter(treatment_state=treatment_state)
+        q.update(treatment_status=equiv)
+        if treatment_state.executed:
+            q.update(input_status=input_status_executed)
+        else:
+            q.update(input_status=input_status_draft)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('archaeological_finds', '0136_treatment_status_refactoring'),
+    ]
+
+    operations = [
+        migrations.RunPython(load_data)
+    ]