Fix Enum Migration to properly deal with constraint changes

Since we're using non-native enums for our model, alembic has some issues knowing what the previous "type" actually was, and not specifying it correctly causes a bad constraint.

alembic/env.py

@@ -29,11 +29,13 @@ parent_dir = Path(__file__).parent.parent
 sys.path.append(parent_dir)
 
 from atst.app import make_config
+
 app_config = make_config()
-config.set_main_option('sqlalchemy.url', app_config['DATABASE_URI'])
+config.set_main_option("sqlalchemy.url", app_config["DATABASE_URI"])
 
 from atst.database import db
 from atst.models import *
+
 target_metadata = Base.metadata
 
 
@@ -51,8 +53,7 @@ def run_migrations_offline():
     """
     url = config.get_main_option("sqlalchemy.url")
     context.configure(
-        url=url, target_metadata=target_metadata, literal_binds=True,
-        compare_type=True
+        url=url, target_metadata=target_metadata, literal_binds=True, compare_type=True
     )
 
     with context.begin_transaction():
@@ -68,19 +69,19 @@ def run_migrations_online():
     """
     connectable = engine_from_config(
         config.get_section(config.config_ini_section),
-        prefix='sqlalchemy.',
-        poolclass=pool.NullPool)
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
 
     with connectable.connect() as connection:
         context.configure(
-            connection=connection,
-            target_metadata=target_metadata,
-            compare_type=True
+            connection=connection, target_metadata=target_metadata, compare_type=True
         )
 
         with context.begin_transaction():
             context.run_migrations()
 
+
 if context.is_offline_mode():
     run_migrations_offline()
 else:

alembic/versions/26319c44a8d5_state_machine_states_extended.py

@@ -10,33 +10,124 @@ import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql
 
 # revision identifiers, used by Alembic.
-revision = '26319c44a8d5' # pragma: allowlist secret
-down_revision = '59973fa17ded' # pragma: allowlist secret
+revision = "26319c44a8d5"  # pragma: allowlist secret
+down_revision = "59973fa17ded"  # pragma: allowlist secret
 branch_labels = None
 depends_on = None
 
 
 def upgrade():
     # ### commands auto generated by Alembic - please adjust! ###
-    op.alter_column('portfolio_state_machines', 'state',
-               existing_type=sa.VARCHAR(length=30),
-               type_=sa.Enum('UNSTARTED', 'STARTING', 'STARTED', 'COMPLETED', 'FAILED', 'TENANT_CREATED', 'TENANT_IN_PROGRESS', 'TENANT_FAILED', 'BILLING_PROFILE_CREATION_CREATED', 'BILLING_PROFILE_CREATION_IN_PROGRESS', 'BILLING_PROFILE_CREATION_FAILED', 'BILLING_PROFILE_VERIFICATION_CREATED', 'BILLING_PROFILE_VERIFICATION_IN_PROGRESS', 'BILLING_PROFILE_VERIFICATION_FAILED', 'BILLING_PROFILE_TENANT_ACCESS_CREATED', 'BILLING_PROFILE_TENANT_ACCESS_IN_PROGRESS', 'BILLING_PROFILE_TENANT_ACCESS_FAILED', 'TASK_ORDER_BILLING_CREATION_CREATED', 'TASK_ORDER_BILLING_CREATION_IN_PROGRESS', 'TASK_ORDER_BILLING_CREATION_FAILED', 'TASK_ORDER_BILLING_VERIFICATION_CREATED', 'TASK_ORDER_BILLING_VERIFICATION_IN_PROGRESS', 'TASK_ORDER_BILLING_VERIFICATION_FAILED', 'BILLING_INSTRUCTION_CREATED', 'BILLING_INSTRUCTION_IN_PROGRESS', 'BILLING_INSTRUCTION_FAILED', name='fsmstates', native_enum=False, create_constraint=False),
-               existing_nullable=False)
-    op.alter_column('portfolios', 'defense_component',
-               existing_type=postgresql.ARRAY(sa.VARCHAR()),
-               type_=sa.String(),
-               existing_nullable=False)
+    op.alter_column(
+        "portfolio_state_machines",
+        "state",
+        existing_type=sa.Enum(
+            "UNSTARTED",
+            "STARTING",
+            "STARTED",
+            "COMPLETED",
+            "FAILED",
+            "TENANT_CREATED",
+            "TENANT_IN_PROGRESS",
+            "TENANT_FAILED",
+            "BILLING_PROFILE_CREATED",
+            "BILLING_PROFILE_IN_PROGRESS",
+            "BILLING_PROFILE_FAILED",
+            "ADMIN_SUBSCRIPTION_CREATED",
+            "ADMIN_SUBSCRIPTION_IN_PROGRESS",
+            "ADMIN_SUBSCRIPTION_FAILED",
+            name="fsmstates",
+            native_enum=False,
+        ),
+        type_=sa.Enum(
+            "UNSTARTED",
+            "STARTING",
+            "STARTED",
+            "COMPLETED",
+            "FAILED",
+            "TENANT_CREATED",
+            "TENANT_IN_PROGRESS",
+            "TENANT_FAILED",
+            "BILLING_PROFILE_CREATION_CREATED",
+            "BILLING_PROFILE_CREATION_IN_PROGRESS",
+            "BILLING_PROFILE_CREATION_FAILED",
+            "BILLING_PROFILE_VERIFICATION_CREATED",
+            "BILLING_PROFILE_VERIFICATION_IN_PROGRESS",
+            "BILLING_PROFILE_VERIFICATION_FAILED",
+            "BILLING_PROFILE_TENANT_ACCESS_CREATED",
+            "BILLING_PROFILE_TENANT_ACCESS_IN_PROGRESS",
+            "BILLING_PROFILE_TENANT_ACCESS_FAILED",
+            "TASK_ORDER_BILLING_CREATION_CREATED",
+            "TASK_ORDER_BILLING_CREATION_IN_PROGRESS",
+            "TASK_ORDER_BILLING_CREATION_FAILED",
+            "TASK_ORDER_BILLING_VERIFICATION_CREATED",
+            "TASK_ORDER_BILLING_VERIFICATION_IN_PROGRESS",
+            "TASK_ORDER_BILLING_VERIFICATION_FAILED",
+            "BILLING_INSTRUCTION_CREATED",
+            "BILLING_INSTRUCTION_IN_PROGRESS",
+            "BILLING_INSTRUCTION_FAILED",
+            name="fsmstates",
+            native_enum=False,
+            create_constraint=False,
+        ),
+        existing_nullable=False,
+    )
     # ### end Alembic commands ###
 
 
 def downgrade():
     # ### commands auto generated by Alembic - please adjust! ###
-    op.alter_column('portfolios', 'defense_component',
-               existing_type=sa.String(),
-               type_=postgresql.ARRAY(sa.VARCHAR()),
-               existing_nullable=False)
-    op.alter_column('portfolio_state_machines', 'state',
-               existing_type=sa.Enum('UNSTARTED', 'STARTING', 'STARTED', 'COMPLETED', 'FAILED', 'TENANT_CREATED', 'TENANT_IN_PROGRESS', 'TENANT_FAILED', 'BILLING_PROFILE_CREATION_CREATED', 'BILLING_PROFILE_CREATION_IN_PROGRESS', 'BILLING_PROFILE_CREATION_FAILED', 'BILLING_PROFILE_VERIFICATION_CREATED', 'BILLING_PROFILE_VERIFICATION_IN_PROGRESS', 'BILLING_PROFILE_VERIFICATION_FAILED', 'BILLING_PROFILE_TENANT_ACCESS_CREATED', 'BILLING_PROFILE_TENANT_ACCESS_IN_PROGRESS', 'BILLING_PROFILE_TENANT_ACCESS_FAILED', 'TASK_ORDER_BILLING_CREATION_CREATED', 'TASK_ORDER_BILLING_CREATION_IN_PROGRESS', 'TASK_ORDER_BILLING_CREATION_FAILED', 'TASK_ORDER_BILLING_VERIFICATION_CREATED', 'TASK_ORDER_BILLING_VERIFICATION_IN_PROGRESS', 'TASK_ORDER_BILLING_VERIFICATION_FAILED', 'BILLING_INSTRUCTION_CREATED', 'BILLING_INSTRUCTION_IN_PROGRESS', 'BILLING_INSTRUCTION_FAILED', name='fsmstates', native_enum=False),
-               type_=sa.VARCHAR(length=30),
-               existing_nullable=False)
+    op.alter_column(
+        "portfolio_state_machines",
+        "state",
+        existing_type=sa.Enum(
+            "UNSTARTED",
+            "STARTING",
+            "STARTED",
+            "COMPLETED",
+            "FAILED",
+            "TENANT_CREATED",
+            "TENANT_IN_PROGRESS",
+            "TENANT_FAILED",
+            "BILLING_PROFILE_CREATION_CREATED",
+            "BILLING_PROFILE_CREATION_IN_PROGRESS",
+            "BILLING_PROFILE_CREATION_FAILED",
+            "BILLING_PROFILE_VERIFICATION_CREATED",
+            "BILLING_PROFILE_VERIFICATION_IN_PROGRESS",
+            "BILLING_PROFILE_VERIFICATION_FAILED",
+            "BILLING_PROFILE_TENANT_ACCESS_CREATED",
+            "BILLING_PROFILE_TENANT_ACCESS_IN_PROGRESS",
+            "BILLING_PROFILE_TENANT_ACCESS_FAILED",
+            "TASK_ORDER_BILLING_CREATION_CREATED",
+            "TASK_ORDER_BILLING_CREATION_IN_PROGRESS",
+            "TASK_ORDER_BILLING_CREATION_FAILED",
+            "TASK_ORDER_BILLING_VERIFICATION_CREATED",
+            "TASK_ORDER_BILLING_VERIFICATION_IN_PROGRESS",
+            "TASK_ORDER_BILLING_VERIFICATION_FAILED",
+            "BILLING_INSTRUCTION_CREATED",
+            "BILLING_INSTRUCTION_IN_PROGRESS",
+            "BILLING_INSTRUCTION_FAILED",
+            name="fsmstates",
+            native_enum=False,
+        ),
+        type_=sa.Enum(
+            "UNSTARTED",
+            "STARTING",
+            "STARTED",
+            "COMPLETED",
+            "FAILED",
+            "TENANT_CREATED",
+            "TENANT_IN_PROGRESS",
+            "TENANT_FAILED",
+            "BILLING_PROFILE_CREATED",
+            "BILLING_PROFILE_IN_PROGRESS",
+            "BILLING_PROFILE_FAILED",
+            "ADMIN_SUBSCRIPTION_CREATED",
+            "ADMIN_SUBSCRIPTION_IN_PROGRESS",
+            "ADMIN_SUBSCRIPTION_FAILED",
+            name="fsmstates",
+            native_enum=False,
+        ),
+        existing_nullable=False,
+    )
     # ### end Alembic commands ###