Merge pull request #337 from cmu-delphi/development

Anotehr round of migrations

Author: dmytrotsko

src/indicatorsets/migrations/0010_originaldataprovider_and_more.py

@@ -1,31 +1,93 @@
-# Generated by Django 5.2.5 on 2026-06-16 22:01
+# Generated by Django 5.2.5 on 2026-06-16 20:35
 
-import django.db.models.deletion
-from django.db import migrations, models
+from django.db import migrations
 
-from indicatorsets.odp_migration import ensure_original_data_provider_schema
 
+INDICATOR_SET_TABLE = "indicatorsets_indicatorset"
+PROVIDER_TABLE = "indicatorsets_originaldataprovider"
+CHAR_COLUMN = "original_data_provider"
+FK_COLUMN = "original_data_provider_id"
 
-def drop_orphan_original_data_provider_table(apps, schema_editor):
-    """Clean up partial/failed deploys before CreateModel runs."""
-    with schema_editor.connection.cursor() as cursor:
-        cursor.execute("SET FOREIGN_KEY_CHECKS=0")
+
+def _get_indicator_set_columns(cursor):
+    cursor.execute(
+        """
+        SELECT COLUMN_NAME, DATA_TYPE
+        FROM information_schema.COLUMNS
+        WHERE TABLE_SCHEMA = DATABASE()
+          AND TABLE_NAME = %s
+        """,
+        [INDICATOR_SET_TABLE],
+    )
+    return {row[0]: row[1] for row in cursor.fetchall()}
+
+
+def _table_exists(cursor, table_name):
+    cursor.execute("SHOW TABLES LIKE %s", [table_name])
+    return cursor.fetchone() is not None
+
+
+def _drop_foreign_keys(cursor):
+    cursor.execute(
+        """
+        SELECT CONSTRAINT_NAME
+        FROM information_schema.KEY_COLUMN_USAGE
+        WHERE TABLE_SCHEMA = DATABASE()
+          AND TABLE_NAME = %s
+          AND COLUMN_NAME = %s
+          AND REFERENCED_TABLE_NAME IS NOT NULL
+        """,
+        [INDICATOR_SET_TABLE, FK_COLUMN],
+    )
+    for (constraint_name,) in cursor.fetchall():
         cursor.execute(
-            """
-            SELECT CONSTRAINT_NAME
-            FROM information_schema.KEY_COLUMN_USAGE
-            WHERE TABLE_SCHEMA = DATABASE()
-              AND TABLE_NAME = 'indicatorsets_indicatorset'
-              AND COLUMN_NAME = 'original_data_provider_id'
-              AND REFERENCED_TABLE_NAME IS NOT NULL
-            """
+            f"ALTER TABLE `{INDICATOR_SET_TABLE}` DROP FOREIGN KEY `{constraint_name}`"
         )
-        for (constraint_name,) in cursor.fetchall():
+
+
+def revert_original_data_provider_schema(apps, schema_editor):
+    """
+    Ensure DB matches CharField original_data_provider with no ODP lookup table.
+
+    Handles fresh DBs, orphan ODP tables from failed deploys, and FK-based schemas.
+    """
+    with schema_editor.connection.cursor() as cursor:
+        if not _table_exists(cursor, INDICATOR_SET_TABLE):
+            return
+
+        columns = _get_indicator_set_columns(cursor)
+        has_char = CHAR_COLUMN in columns
+        has_fk = FK_COLUMN in columns
+        has_provider_table = _table_exists(cursor, PROVIDER_TABLE)
+
+        cursor.execute("SET FOREIGN_KEY_CHECKS=0")
+
+        if has_fk:
+            if has_provider_table:
+                if not has_char:
+                    cursor.execute(
+                        f"ALTER TABLE `{INDICATOR_SET_TABLE}` "
+                        f"ADD COLUMN `{CHAR_COLUMN}` varchar(255) NOT NULL DEFAULT ''"
+                    )
+                    has_char = True
+
+                cursor.execute(
+                    f"""
+                    UPDATE `{INDICATOR_SET_TABLE}` iset
+                    LEFT JOIN `{PROVIDER_TABLE}` odp
+                      ON iset.`{FK_COLUMN}` = odp.`id`
+                    SET iset.`{CHAR_COLUMN}` = COALESCE(odp.`name`, '')
+                    """
+                )
+
+            _drop_foreign_keys(cursor)
             cursor.execute(
-                f"ALTER TABLE `indicatorsets_indicatorset` "
-                f"DROP FOREIGN KEY `{constraint_name}`"
+                f"ALTER TABLE `{INDICATOR_SET_TABLE}` DROP COLUMN `{FK_COLUMN}`"
             )
-        cursor.execute("DROP TABLE IF EXISTS `indicatorsets_originaldataprovider`")
+
+        if has_provider_table:
+            cursor.execute(f"DROP TABLE IF EXISTS `{PROVIDER_TABLE}`")
+
         cursor.execute("SET FOREIGN_KEY_CHECKS=1")
 
 
@@ -37,60 +99,7 @@ class Migration(migrations.Migration):
 
     operations = [
         migrations.RunPython(
-            drop_orphan_original_data_provider_table,
+            revert_original_data_provider_schema,
             migrations.RunPython.noop,
         ),
-        migrations.CreateModel(
-            name="OriginalDataProvider",
-            fields=[
-                (
-                    "id",
-                    models.BigAutoField(
-                        auto_created=True,
-                        primary_key=True,
-                        serialize=False,
-                        verbose_name="ID",
-                    ),
-                ),
-                (
-                    "name",
-                    models.CharField(max_length=255, unique=True, verbose_name="Name"),
-                ),
-                (
-                    "group",
-                    models.CharField(
-                        choices=[
-                            ("us_government", "U.S. Government"),
-                            ("us_states", "U.S. States"),
-                            ("individual", "Individual providers"),
-                        ],
-                        "constraints": [
-                            models.UniqueConstraint(
-                                fields=("name", "group"),
-                                name="unique_data_provider_name_group",
-                            )
-                        ],
-                    },
-                ),
-                migrations.AlterField(
-                    model_name="indicatorset",
-                    name="original_data_provider",
-                    field=models.ForeignKey(
-                        blank=True,
-                        help_text="Original data provider of the Indicator Set",
-                        null=True,
-                        on_delete=django.db.models.deletion.SET_NULL,
-                        related_name="indicator_sets",
-                        to="indicatorsets.originaldataprovider",
-                        verbose_name="Original Data Provider",
-                    ),
-                ),
-            ],
-            database_operations=[
-                migrations.RunPython(
-                    ensure_original_data_provider_schema,
-                    migrations.RunPython.noop,
-                ),
-            ],
-        ),
     ]