Alembic: how to add a new enum type in migration script without resorting to `alembic.op.execute`

Question

I'm trying to create a new migration on Alembic, that adds a new column of a new Enum type to an existing table. But I'm getting an error that I thought Alembic would've automatically handled.

I'm using Postgres 9.6.6, Alembic 0.9.10, and SQLAlchemy 1.2.9.

My script has this in it:

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "task",
        sa.Column(
            "task_type",
            sa.Enum(
                "train", "predict", "refresh", name="tasktypeenum"
            ),
            nullable=True,
        ),
    )
    # ### end Alembic commands ###

Here's the end of the stack trace when I try and apply this migration script (the rest is just from SQLAlchemy):

sqlalchemy.exc.ProgrammingError: (psycopg2.ProgrammingError) type "tasktypeenum" does not exist
LINE 1: ...task ADD COLUMN task_type tasktypeen...
                                     ^
[SQL: 'ALTER TABLE task ADD COLUMN task_type tasktypeenum'] (Background on this error at: http://sqlalche.me/e/f405)

However, we have similar statements in old migration files which work fine. The only difference I see is they're nested inside of a op.create_table.

I'm aware I could get around this by doing op.execute("CREATE TYPE ..."), but I'm wondering if there's a solution that uses Alembic's features.

Thanks for any guidance!

Answer 1

You can do something like this:

task_type_enum = sa.Enum("train", "predict", "refresh", name="tasktypeenum")

def upgrade():
    task_type_enum.create(op.get_bind())
    op.add_column("task", sa.Column("task_type", task_type_enum, nullable=True))

def downgrade():
    op.drop_column("task", "task_type")
    task_type_enum.drop(op.get_bind())