Add / remove a value of PostgreSQL ENUM type in alembic

Question

I have a trouble with altering existing postgresql.ENUM column with SQLAlchemy and Alembic.

I want to add / remove a value to postgresql.ENUM type column in alembic.

In specific, current enum type is created by the following two alembic revisions:

# revision 1
def upgrade():
    op.create_table('kernels',
        sa.Column('status', sa.String(), nullable=True),
        ...
    )


# revision 2
kernelstatus_choices = (
    'PREPARING', 'BUILDING', 'RUNNING',
    'RESTARTING', 'RESIZING', 'SUSPENDED',
    'TERMINATING', 'TERMINATED', 'ERROR',
)
kernelstatus = postgresql.ENUM(
    *kernelstatus_choices,
    name='kernelstatus')

def upgrade():
    op.alter_column('kernels', column_name='status', 
                    type_=sa.Enum(*kernelstatus_choices, name='kernelstatus'),
                    postgresql_using='status::kernelstatus')

Now, I want to add 'PENDING' status to kernelstatus type. So I implemented like below, by referencing some articles.

prev_kernelstatus_choices = (
    'PREPARING', 'BUILDING', 'RUNNING',
    'RESTARTING', 'RESIZING', 'SUSPENDED',
    'TERMINATING', 'TERMINATED', 'ERROR',
)

prev_kernelstatus = postgresql.ENUM(
    *prev_kernelstatus_choices,
    name='kernelstatus')

curr_kernelstatus_choices = ('PENDING',) + prev_kernelstatus_choices

curr_kernelstatus = postgresql.ENUM(
    *curr_kernelstatus_choices,
    name='kernelstatus')

def upgrade():
    op.execute('ALTER TYPE kernelstatus RENAME TO kernelstatus_old;')
    curr_kernelstatus.create(op.get_bind())
    op.alter_column('kernels', column_name='status', 
                    type_=sa.Enum(*curr_kernelstatus_choices, name='kernelstatus'),
                    postgresql_using='status::text::kernelstatus')
    op.execute('DROP TYPE kernelstatus_old;')

But it keeps generating the following error:

sqlalchemy.exc.ProgrammingError: (psycopg2.ProgrammingError) operator does not exist: kernelstatus <> kernelstatus_old
HINT:  No operator matches the given name and argument type(s). You might need to add explicit type casts.
[SQL: 'ALTER TABLE kernels ALTER COLUMN status TYPE kernelstatus USING status::text::kernelstatus']

I already tried a solution with adding value to enum type, but this does not work in with Alembic since each Alembic revision runs in a transaction and ALTER TYPE statement cannot run in a transaction. Also, there should be a code for downgrade() and there is no statement for removing a value from enum type in PostgreSQL, so just adding a value to the enum type cannot be the ultimate solution in my case.

Could somebody give me a help?

Answer 1

I figured out the root cause of this case (I've done collaboration with the questioner): since the kernels table has a unique constraint that checks the uniqueness of (the sess_id column, the status column of the enum type kernelstatus when its value is 'TERMINATED').

Converting the enum type of the status column requires an explicit definition of <> SQL operator to keep this constraint check consistent across the conversion. The problem was that PostgreSQL's error message just said "operator does not exist" rather than "operator must be defined to keep constraint during conversion".

So, I could write a working migration, by temporarily adding a comparison operator implemented by converting both old/new enum values into strings:

from alembic import op
import sqlalchemy as sa
import textwrap
from sqlalchemy.dialects import postgresql

kernelstatus_new_values = [
    'PENDING',     # added
    'PREPARING',
    ...
]
kernelstatus_new = postgresql.ENUM(*kernelstatus_new_values, name='kernelstatus')

kernelstatus_old_values = [
    'PREPARING',
    ...
]
kernelstatus_old = postgresql.ENUM(*kernelstatus_old_values, name='kernelstatus')

def upgrade():
    conn = op.get_bind()
    sessionresult.create(conn)
    sessiontypes.create(conn)
    conn.execute('ALTER TYPE kernelstatus RENAME TO kernelstatus_old;')
    kernelstatus_new.create(conn)
    conn.execute(textwrap.dedent('''\
    CREATE FUNCTION new_old_compare(
        new_enum_val kernelstatus, old_enum_val kernelstatus_old
    )
        RETURNS boolean AS $$
            SELECT new_enum_val::text <> old_enum_val::text;
        $$ LANGUAGE SQL IMMUTABLE;

    CREATE OPERATOR <> (
        leftarg = kernelstatus,
        rightarg = kernelstatus_old,
        procedure = new_old_compare
    );
    '''))
    op.alter_column(
        table_name='kernels',
        column_name='status',
        type_=kernelstatus_new,
        postgresql_using='status::text::kernelstatus',
    )
    conn.execute(textwrap.dedent('''\
    DROP FUNCTION new_old_compare(
        new_enum_val kernelstatus, old_enum_val kernelstatus_old
    ) CASCADE;
    DROP TYPE kernelstatus_old;
    '''))

    ...  # the rest of migration