在 alembic 中添加/删除 PostgreSQL ENUM 类型的值

Question

我无法使用 SQLAlchemy 和 Alembic 更改现有的 postgresql.ENUM 列。

我想向 alembic 中的 postgresql.ENUM 类型列添加/删除一个值。

具体来说，当前枚举类型是由以下两个 alembic 修订创建的：

# 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')

现在，我想将'PENDING' 状态添加到kernelstatus 类型。所以我通过引用somearticles来实现如下所示。

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;')

但它不断产生以下错误：

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']

我已经尝试过使用adding value to enum type 的解决方案，但这不适用于Alembic，因为每个Alembic 修订版都在事务中运行，而ALTER TYPE 语句不能在事务中运行。另外，downgrade() 应该有一个代码，并且在 PostgreSQL 中没有从枚举类型中删除值的语句，因此在我的情况下，仅向枚举类型添加一个值并不是最终的解决方案。

有人可以帮我吗？

Answer 1

我找到了这个案例的根本原因（我已经与提问者合作）：因为kernels 表有一个唯一性约束来检查（sess_id 列、status 列的唯一性枚举类型kernelstatus，当其值为'TERMINATED'时。

转换status 列的枚举类型需要显式定义<> SQL 运算符，以使此约束检查在整个转换过程中保持一致。 问题是 PostgreSQL 的错误信息只是说“操作符不存在”而不是“必须定义操作符以在转换期间保持约束”。

所以，我可以编写一个有效的迁移，通过临时添加一个比较运算符，通过将旧/新枚举值转换为字符串来实现：

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