fix migration

migrations/versions/0418_user_state_enum.py

@@ -5,38 +5,152 @@ Revises: 0417_change_total_message_limit
 Create Date: 2025-08-28 12:34:32.857422
 
 """
-import sqlalchemy as sa
-from alembic import op
-from sqlalchemy.dialects import postgresql
-
 
 revision = "0418_user_state_enum"
 down_revision = "0417_change_total_message_limit"
 
+from contextlib import contextmanager
+from enum import Enum
+from typing import Iterator, TypedDict
 
-user_state_enum = postgresql.ENUM(
-    "active", "pending", "inactive", name="user_states", create_type=False
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+from app.enums import (
+    UserState,
 )
 
 
-def upgrade():
-    user_state_enum.create(op.get_bind(), checkfirst=True)
-    op.alter_column(
-        "users",
-        "state",
-        existing_type=sa.String(),
-        type_=user_state_enum,
-        existing_nullable=False,
+class EnumValues(TypedDict):
+    values: list[str]
+    name: str
+
+
+_enum_params: dict[Enum, EnumValues] = {
+    UserState: {
+        "values": ["active", "inactive", "pending"],
+        "name": "user_states",
+    },
+}
+
+
+def enum_create(values: list[str], name: str) -> None:
+    enum_db_type = postgresql.ENUM(*values, name=name)
+    enum_db_type.create(op.get_bind())
+
+
+def enum_drop(values: list[str], name: str) -> None:
+    enum_db_type = postgresql.ENUM(*values, name=name)
+    enum_db_type.drop(op.get_bind())
+
+
+def enum_using(column_name: str, enum: Enum) -> str:
+    return f"{column_name}::text::{_enum_params[enum]['name']}"
+
+
+def enum_type(enum: Enum) -> sa.Enum:
+    return sa.Enum(
+        *_enum_params[enum]["values"],
+        name=_enum_params[enum]["name"],
+        values_callable=(lambda x: [e.value for e in x]),
     )
 
 
+@contextmanager
+def view_handler() -> Iterator[None]:
+    op.execute("DROP VIEW notifications_all_time_view")
+
+    yield
+
+    op.execute(
+        """
+        CREATE VIEW notifications_all_time_view AS
+        (
+            SELECT
+                id,
+                job_id,
+                job_row_number,
+                service_id,
+                template_id,
+                template_version,
+                api_key_id,
+                key_type,
+                billable_units,
+                notification_type,
+                created_at,
+                sent_at,
+                sent_by,
+                updated_at,
+                notification_status,
+                reference,
+                client_reference,
+                international,
+                phone_prefix,
+                rate_multiplier,
+                created_by_id,
+                document_download_count
+            FROM notifications
+        ) UNION
+        (
+            SELECT
+                id,
+                job_id,
+                job_row_number,
+                service_id,
+                template_id,
+                template_version,
+                api_key_id,
+                key_type,
+                billable_units,
+                notification_type,
+                created_at,
+                sent_at,
+                sent_by,
+                updated_at,
+                notification_status,
+                reference,
+                client_reference,
+                international,
+                phone_prefix,
+                rate_multiplier,
+                created_by_id,
+                document_download_count
+            FROM notification_history
+        )
+    """
+    )
+
+
+def upgrade():
+    with view_handler():
+
+        for enum_data in _enum_params.values():
+            enum_create(**enum_data)
+
+        # alter existing columns to use new enums
+        op.alter_column(
+            "users",
+            "state",
+            existing_type=sa.VARCHAR(length=255),
+            type_=enum_type(UserState),
+            existing_nullable=True,
+            postgresql_using=enum_using("user_states", UserState),
+        )
+
+
 def downgrade():
-    op.alter_column(
-        "users",
-        "state",
-        existing_type=user_state_enum,
-        type_=sa.String,
-        existing_nullable=False,
-    )
+    with view_handler():
+        # Create old enum types.
+        # Alter columns back
 
-    user_state_enum.drop(op.get_bind(), checkfirst=True)
+        op.alter_column(
+            "users",
+            "state",
+            existing_type=enum_type(UserState),
+            type_=sa.VARCHAR(length=255),
+            existing_nullable=True,
+        )
+
+        for enum_data in _enum_params.values():
+            enum_drop(**enum_data)