From 0cf9ef5b9a7ba070d0aaafeaf54d98d300d4dfd5 Mon Sep 17 00:00:00 2001
From: AlexCLeduc <ledualex@gmail.com>
Date: Wed, 20 Sep 2023 16:55:43 -0400
Subject: [PATCH] add custom migration class and prototype example

---
 .../migrations/0008_load_period_fixtures.py   | 27 ++++++++++++++++
 server/server/model_util.py                   | 31 ++++++++++++++++++-
 2 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 server/cpho/migrations/0008_load_period_fixtures.py

diff --git a/server/cpho/migrations/0008_load_period_fixtures.py b/server/cpho/migrations/0008_load_period_fixtures.py
new file mode 100644
index 00000000..6aa7c70a
--- /dev/null
+++ b/server/cpho/migrations/0008_load_period_fixtures.py
@@ -0,0 +1,27 @@
+# Generated by Django 4.1.9 on 2023-09-20 17:59
+
+from django.db import migrations
+from django.core.management import call_command
+
+from server.model_util import DevProdTestMigration
+
+
+def up(apps, schema_editor):
+    call_command("loaddata", "cpho/fixtures/periods.yaml")
+
+
+def down(apps, schema_editor):
+    Period = apps.get_model("cpho", "Period")
+    Period.objects.all().delete()
+
+
+class Migration(DevProdTestMigration):
+    apply_in_dev = True
+    apply_in_prod = True
+    apply_in_tests = False
+
+    dependencies = [
+        ("cpho", "0007_alter_dimensionvalue_name_en_and_more"),
+    ]
+
+    operations = [migrations.RunPython(up, down)]
diff --git a/server/server/model_util.py b/server/server/model_util.py
index 3864bf33..cb6af9c4 100644
--- a/server/server/model_util.py
+++ b/server/server/model_util.py
@@ -1,7 +1,7 @@
 import importlib
 
 from django.conf import settings
-from django.db import models
+from django.db import migrations, models
 from django.utils import timezone
 
 from versionator import VersionModel
@@ -88,3 +88,32 @@ def decorator(live_model):
 track_versions_with_editor_and_submission = create_history_decorator(
     ApprovableCustomVersionModelWithEditor
 )
+
+
+class EnvironmentDependentMigration(migrations.Migration):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if not self.should_run():
+            self.operations = []
+
+    def should_skip(self):
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class DevProdTestMigration(EnvironmentDependentMigration):
+    def should_run(self):
+        attrs_to_check = ["apply_in_tests", "apply_in_prod", "apply_in_dev"]
+        if any([not hasattr(self.__class__, attr) for attr in attrs_to_check]):
+            raise NotImplementedError(
+                "DevProdTestMigration must specify apply_in_tests, apply_in_prod and apply_in_dev bool attributes"
+            )
+
+        if settings.IS_RUNNING_TESTS:
+            return self.__class__.apply_in_tests is True
+
+        if settings.IS_DEV:
+            return self.__class__.apply_in_dev is True
+
+        # TODO: have a more explicit check for prov and raise if none of dev/prod/test are true
+        return self.__class__.apply_in_prod is True