在 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 类型。所以我通过引用 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;')

但它不断产生以下错误：

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 修订运行s 和 ALTER TYPE 语句不能运行在交易中。另外，downgrade()应该有一个代码，并且在PostgreSQL中没有从枚举类型中删除一个值的语句，所以在我的情况下，仅仅向枚举类型添加一个值并不是最终的解决方案。

有人可以帮我吗？

Answer 1

我找出了这种情况的根本原因（我已经与提问者合作）：因为 kernels table 有一个唯一的约束来检查（sess_id列，枚举类型kernelstatus的status列，当它的值为'TERMINATED').

转换 status 列的枚举类型需要 <> SQL 运算符的显式定义，以使此约束检查在整个转换过程中保持一致。问题是 PostgreSQL 的错误消息只是说 "operator does not exist" 而不是 "operator must be defined to keep constraint during conversion".

因此，我可以通过临时添加一个比较运算符来编写工作迁移，该比较运算符通过将两个 old/new 枚举值转换为字符串来实现：

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

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

Add / remove a value of PostgreSQL ENUM type in alembic

postgresql

enums

sqlalchemy

alembic