More migration fun.

Signed-off-by: Cliff Hill <Clifford.hill@gsa.gov>

app/models.py

@@ -14,7 +14,7 @@ from notifications_utils.recipients import (
 )
 from notifications_utils.template import PlainTextEmailTemplate, SMSMessageTemplate
 from sqlalchemy import CheckConstraint, Index, UniqueConstraint
-from sqlalchemy.dialects.postgresql import JSON, JSONB, UUID
+from sqlalchemy.dialects.postgresql import JSON, JSONB, UUID, ENUM
 from sqlalchemy.ext.associationproxy import association_proxy
 from sqlalchemy.ext.declarative import declared_attr
 from sqlalchemy.orm import validates
@@ -53,120 +53,35 @@ def filter_null_value_fields(obj):
     return dict(filter(lambda x: x[1] is not None, obj.items()))
 
 
-def enum_values(enum: Enum) -> list[str]:
-    """
-    Helper function used to persist enum values to the database rather than names.
-
-    See Also:
-        https://docs.sqlalchemy.org/en/14/core/type_basics.html#sqlalchemy.types.Enum
-
-    Notes:
-        In order to persist the values and not the names, the Enum.values_callable
-        parameter may be used. The value of this parameter is a user-supplied callable,
-        which is intended to be used with a PEP-435-compliant enumerated class and
-        returns a list of string values to be persisted. For a simple enumeration that
-        uses string values, a callable such as lambda x: [e.value for e in x] is
-        sufficient.
-    """
-    return [i.value for i in enum]  # type: ignore[attr-defined]
-
-
-# This has the standard definition for all of the enum columns used throughout the models.
-# This is used in the enum_column() function below.
-_enum_column_types = {
-    AuthType: db.Enum(  # the key should be the Enum class to use.
-        AuthType,  # The Enum class to use.
-        name="auth_types",  # The name of the type defined in PostgreSQL - should be plural
-        values_callable=enum_values,  # Every entry in this dict should contain this, note above.
-    ),
-    BrandType: db.Enum(
-        BrandType,
-        name="brand_types",
-        values_callable=enum_values,
-    ),
-    OrganizationType: db.Enum(
-        OrganizationType,
-        name="organization_types",
-        values_callable=enum_values,
-    ),
-    ServicePermissionType: db.Enum(
-        ServicePermissionType,
-        name="service_permission_types",
-        values_callable=enum_values,
-    ),
-    RecipientType: db.Enum(
-        RecipientType,
-        name="recipient_types",
-        values_callable=enum_values,
-    ),
-    CallbackType: db.Enum(
-        CallbackType,
-        name="callback_types",
-        values_callable=enum_values,
-    ),
-    KeyType: db.Enum(
-        KeyType,
-        name="key_types",
-        values_callable=enum_values,
-    ),
-    TemplateType: db.Enum(
-        TemplateType,
-        name="template_types",
-        values_callable=enum_values,
-    ),
-    TemplateProcessType: db.Enum(
-        TemplateProcessType,
-        name="template_process_types",
-        values_callable=enum_values,
-    ),
-    NotificationType: db.Enum(
-        NotificationType,
-        name="notification_types",
-        values_callable=enum_values,
-    ),
-    JobStatus: db.Enum(
-        JobStatus,
-        name="job_statuses",
-        values_callable=enum_values,
-    ),
-    CodeType: db.Enum(
-        CodeType,
-        name="code_types",
-        values_callable=enum_values,
-    ),
-    NotificationStatus: db.Enum(
-        NotificationStatus,
-        name="notify_statuses",
-        values_callable=enum_values,
-    ),
-    InvitedUserStatus: db.Enum(
-        InvitedUserStatus,
-        name="invited_user_statuses",
-        values_callable=enum_values,
-    ),
-    PermissionType: db.Enum(
-        PermissionType,
-        name="permission_types",
-        values_callable=enum_values,
-    ),
-    AgreementType: db.Enum(
-        AgreementType,
-        name="agreement_types",
-        values_callable=enum_values,
-    ),
-    AgreementStatus: db.Enum(
-        AgreementStatus,
-        name="agreement_statuses",
-        values_callable=enum_values,
-    ),
+_enum_column_names = {
+    AuthType: "auth_types",
+    BrandType: "brand_types",
+    OrganizationType: "organization_types",
+    ServicePermissionType: "service_permission_types",
+    RecipientType: "recipient_types",
+    CallbackType: "callback_types",
+    KeyType: "key_types",
+    TemplateType: "template_types",
+    TemplateProcessType: "template_process_types",
+    NotificationType: "notification_types",
+    JobStatus: "job_statuses",
+    CodeType: "code_types",
+    NotificationStatus: "notify_statuses",
+    InvitedUserStatus: "invited_user_statuses",
+    PermissionType: "permission_types",
+    AgreementType: "agreement_types",
+    AgreementStatus: "agreement_statuses",
 }
 
 
-def enum_column(enum_type, name=None, **kwargs):
-    if name is None:
-        return db.Column(_enum_column_types[enum_type], **kwargs)
-    else:
-        return db.Column(name, _enum_column_types[enum_type], **kwargs)
+def enum_column(enum_type, **kwargs):
+    return db.Column(
+        db.Enum(
+            *[i.value for i in enum_type],
+            name=_enum_column_names[enum_type]
+        ),
+        **kwargs,
+    )
 
 
 class HistoryModel:
@@ -1591,6 +1506,7 @@ class Notification(db.Model):
         name="notification_status",
         nullable=True,
         default=NotificationStatus.CREATED,
+        key="status",
     )
     reference = db.Column(db.String, nullable=True, index=True)
     client_reference = db.Column(db.String, index=True, nullable=True)
@@ -1867,6 +1783,7 @@ class NotificationHistory(db.Model, HistoryModel):
         name="notification_status",
         nullable=True,
         default=NotificationStatus.CREATED,
+        key="status",
     )
     reference = db.Column(db.String, nullable=True, index=True)
     client_reference = db.Column(db.String, nullable=True)

migrations/versions/0410_enums_for_everything.py

@@ -118,6 +118,19 @@ _enum_params: dict[Enum, EnumValues] = {
     CodeType: {"values": ["email", "sms"], "name": "code_types"},
 }
 
+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}::{_enum_params[enum]['name']}"
+
 
 def enum_type(enum: Enum) -> sa.Enum:
     return sa.Enum(*_enum_params[enum]["values"], name=_enum_params[enum]["name"])
@@ -188,6 +201,9 @@ def upgrade():
     op.drop_table("job_status")
     op.drop_table("service_permission_types")
 
+    for enum_data in _enum_params.values():
+        enum_create(**enum_data)
+
     # alter existing columns to use new enums
     op.alter_column(
         "api_keys",
@@ -195,6 +211,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(KeyType),
         existing_nullable=False,
+        postgresql_using=enum_using("key_type", KeyType),
     )
     op.alter_column(
         "api_keys_history",
@@ -202,6 +219,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(KeyType),
         existing_nullable=False,
+        postgresql_using=enum_using("key_type", KeyType),
     )
     op.alter_column(
         "email_branding",
@@ -209,6 +227,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(BrandType),
         existing_nullable=False,
+        postgresql_using=enum_using("brand_type", BrandType),
     )
     op.alter_column(
         "invited_organization_users",
@@ -216,6 +235,7 @@ def upgrade():
         existing_type=sa.VARCHAR(),
         type_=enum_type(InvitedUserStatus),
         existing_nullable=False,
+        postgresql_using=enum_using("status", InvitedUserStatus),
     )
     op.alter_column(
         "invited_users",
@@ -229,6 +249,7 @@ def upgrade():
         ),
         type_=enum_type(InvitedUserStatus),
         existing_nullable=False,
+        postgresql_using=enum_using("status", InvitedUserStatus),
     )
     op.alter_column(
         "invited_users",
@@ -237,6 +258,7 @@ def upgrade():
         type_=enum_type(AuthType),
         existing_nullable=False,
         existing_server_default=sa.text("'sms_auth'::character varying"),
+        postgresql_using=enum_using("auth_type", AuthType),
     )
     op.alter_column(
         "jobs",
@@ -244,6 +266,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(JobStatus),
         existing_nullable=False,
+        postgresql_using=enum_using("job_status", JobStatus),
     )
     op.alter_column(
         "notification_history",
@@ -251,6 +274,7 @@ def upgrade():
         existing_type=sa.TEXT(),
         type_=enum_type(NotificationStatus),
         existing_nullable=True,
+        postgresql_using=enum_using("notification_status", NotificationStatus),
     )
     op.alter_column(
         "notification_history",
@@ -258,6 +282,7 @@ def upgrade():
         existing_type=sa.VARCHAR(),
         type_=enum_type(KeyType),
         existing_nullable=False,
+        postgresql_using=enum_using("key_type", KeyType),
     )
     op.alter_column(
         "notification_history",
@@ -267,6 +292,7 @@ def upgrade():
         ),
         type_=enum_type(NotificationType),
         existing_nullable=False,
+        postgresql_using=enum_using("notification_type", NotificationType),
     )
     op.alter_column(
         "notifications",
@@ -274,6 +300,7 @@ def upgrade():
         existing_type=sa.TEXT(),
         type_=enum_type(NotificationStatus),
         existing_nullable=True,
+        postgresql_using=enum_using("notification_status", NotificationStatus),
     )
     op.alter_column(
         "notifications",
@@ -281,6 +308,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(KeyType),
         existing_nullable=False,
+        postgresql_using=enum_using("key_type", KeyType),
     )
     op.alter_column(
         "notifications",
@@ -290,9 +318,7 @@ def upgrade():
         ),
         type_=enum_type(NotificationType),
         existing_nullable=False,
-    )
-    op.alter_column(
-        "notifications", "international", existing_type=sa.BOOLEAN(), nullable=False
+        postgresql_using=enum_using("notification_type", NotificationType),
     )
     op.alter_column(
         "organization",
@@ -300,6 +326,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(OrganizationType),
         existing_nullable=True,
+        postgresql_using=enum_using("organization_type", OrganizationType),
     )
     op.alter_column(
         "provider_details",
@@ -309,6 +336,7 @@ def upgrade():
         ),
         type_=enum_type(NotificationType),
         existing_nullable=False,
+        postgresql_using=enum_using("notification_type", NotificationType),
     )
     op.alter_column(
         "provider_details_history",
@@ -318,6 +346,7 @@ def upgrade():
         ),
         type_=enum_type(NotificationType),
         existing_nullable=False,
+        postgresql_using=enum_using("notification_type", NotificationType),
     )
     op.alter_column(
         "rates",
@@ -327,6 +356,7 @@ def upgrade():
         ),
         type_=enum_type(NotificationType),
         existing_nullable=False,
+        postgresql_using=enum_using("notification_type", NotificationType),
     )
     op.alter_column(
         "service_callback_api",
@@ -334,6 +364,7 @@ def upgrade():
         existing_type=sa.VARCHAR(),
         type_=enum_type(CallbackType),
         existing_nullable=True,
+        postgresql_using=enum_using("callback_type", CallbackType),
     )
     op.alter_column(
         "service_callback_api_history",
@@ -341,6 +372,7 @@ def upgrade():
         existing_type=sa.VARCHAR(),
         type_=enum_type(CallbackType),
         existing_nullable=True,
+        postgresql_using=enum_using("callback_type", CallbackType),
     )
     op.alter_column(
         "service_data_retention",
@@ -350,6 +382,7 @@ def upgrade():
         ),
         type_=enum_type(NotificationType),
         existing_nullable=False,
+        postgresql_using=enum_using("notification_type", NotificationType),
     )
     op.alter_column(
         "service_permissions",
@@ -357,6 +390,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(ServicePermissionType),
         existing_nullable=False,
+        postgresql_using=enum_using("permission", ServicePermissionType),
     )
     op.alter_column(
         "service_whitelist",
@@ -364,6 +398,7 @@ def upgrade():
         existing_type=postgresql.ENUM("mobile", "email", name="recipient_type"),
         type_=enum_type(RecipientType),
         existing_nullable=False,
+        postgresql_using=enum_using("recipient_type", RecipientType),
     )
     op.alter_column(
         "services",
@@ -371,6 +406,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(OrganizationType),
         existing_nullable=True,
+        postgresql_using=enum_using("organization_type", OrganizationType),
     )
     op.alter_column(
         "services_history",
@@ -378,6 +414,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(OrganizationType),
         existing_nullable=True,
+        postgresql_using=enum_using("organization_type", OrganizationType),
     )
     op.alter_column(
         "templates",
@@ -387,6 +424,7 @@ def upgrade():
         ),
         type_=enum_type(TemplateType),
         existing_nullable=False,
+        postgresql_using=enum_using("template_type", TemplateType),
     )
     op.alter_column(
         "templates",
@@ -394,6 +432,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(TemplateProcessType),
         existing_nullable=False,
+        postgresql_using=enum_using("process_type", TemplateProcessType),
     )
     op.alter_column(
         "templates_history",
@@ -403,6 +442,7 @@ def upgrade():
         ),
         type_=enum_type(TemplateType),
         existing_nullable=False,
+        postgresql_using=enum_using("template_type", TemplateType),
     )
     op.alter_column(
         "templates_history",
@@ -410,6 +450,7 @@ def upgrade():
         existing_type=sa.VARCHAR(length=255),
         type_=enum_type(TemplateProcessType),
         existing_nullable=False,
+        postgresql_using=enum_using("process_type", TemplateProcessType),
     )
     op.alter_column(
         "users",
@@ -418,6 +459,7 @@ def upgrade():
         type_=enum_type(AuthType),
         existing_nullable=False,
         existing_server_default=sa.text("'sms_auth'::character varying"),
+        postgresql_using=enum_using("auth_type", AuthType),
     )
     op.alter_column(
         "verify_codes",
@@ -425,6 +467,7 @@ def upgrade():
         existing_type=postgresql.ENUM("email", "sms", name="verify_code_types"),
         type_=enum_type(CodeType),
         existing_nullable=False,
+        postgresql_using=enum_using("code_type", CodeType),
     )
 
     # Recreate composite indexes
@@ -458,6 +501,7 @@ def downgrade():
         existing_type=enum_type(CodeType),
         type_=postgresql.ENUM("email", "sms", name="verify_code_types"),
         existing_nullable=False,
+        postgresql_using="code_type::verify_code_types",
     )
     op.alter_column(
         "users",
@@ -482,6 +526,7 @@ def downgrade():
             "sms", "email", "letter", "broadcast", name="template_type"
         ),
         existing_nullable=False,
+        postgresql_using="template_type::template_type",
     )
     op.alter_column(
         "templates",
@@ -498,6 +543,7 @@ def downgrade():
             "sms", "email", "letter", "broadcast", name="template_type"
         ),
         existing_nullable=False,
+        postgresql_using="template_type::template_type",
     )
     op.alter_column(
         "services_history",
@@ -519,6 +565,7 @@ def downgrade():
         existing_type=enum_type(RecipientType),
         type_=postgresql.ENUM("mobile", "email", name="recipient_type"),
         existing_nullable=False,
+        postgresql_using="recipient_type::recipient_type",
     )
     op.alter_column(
         "service_permissions",
@@ -533,6 +580,7 @@ def downgrade():
         existing_type=enum_type(NotificationType),
         type_=postgresql.ENUM("email", "sms", "letter", name="notification_type"),
         existing_nullable=False,
+        postgresql_using="notification_type::notification_type",
     )
     op.alter_column(
         "service_callback_api_history",
@@ -554,6 +602,7 @@ def downgrade():
         existing_type=enum_type(NotificationType),
         type_=postgresql.ENUM("email", "sms", "letter", name="notification_type"),
         existing_nullable=False,
+        postgresql_using="notification_type::notification_type",
     )
     op.alter_column(
         "provider_details_history",
@@ -561,6 +610,7 @@ def downgrade():
         existing_type=enum_type(NotificationType),
         type_=postgresql.ENUM("email", "sms", "letter", name="notification_type"),
         existing_nullable=False,
+        postgresql_using="notification_type::notification_type",
     )
     op.alter_column(
         "provider_details",
@@ -568,6 +618,7 @@ def downgrade():
         existing_type=enum_type(NotificationType),
         type_=postgresql.ENUM("email", "sms", "letter", name="notification_type"),
         existing_nullable=False,
+        postgresql_using="notification_type::notification_type",
     )
     op.alter_column(
         "organization",
@@ -589,6 +640,7 @@ def downgrade():
         existing_type=enum_type(NotificationType),
         type_=postgresql.ENUM("email", "sms", "letter", name="notification_type"),
         existing_nullable=False,
+        postgresql_using="notification_type::notification_type",
     )
     op.alter_column(
         "notifications",
@@ -610,6 +662,7 @@ def downgrade():
         existing_type=enum_type(NotificationType),
         type_=postgresql.ENUM("email", "sms", "letter", name="notification_type"),
         existing_nullable=False,
+        postgresql_using="notification_type::notification_type",
     )
     op.alter_column(
         "notification_history",
@@ -645,6 +698,7 @@ def downgrade():
             name="invited_users_status_types",
         ),
         existing_nullable=False,
+        postgresql_using="status::invited_user_status_types",
     )
     op.alter_column(
         "invited_organization_users",
@@ -675,6 +729,9 @@ def downgrade():
         existing_nullable=False,
     )
 
+    for enum_data in _enum_params.values():
+        enum_drop(**enum_data)
+
     # Recreate helper tables
     service_permission_types = op.create_table(
         "service_permission_types",