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

18 浏览
0 Comments

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

我在使用SQLAlchemy和Alembic时遇到了一个问题，即更改现有的postgresql.ENUM列。\n我想在alembic中向postgresql.ENUM类型的列添加/删除一个值。\n具体而言，当前的枚举类型是通过以下两个alembic修订创建的：\n

# 修订1
def 升级():
    op.create_table('kernels',
        sa.Column('status', sa.String(), nullable=True),
        ...
    )
# 修订2
kernelstatus_choices = (
    'PREPARING', 'BUILDING', 'RUNNING',
    'RESTARTING', 'RESIZING', 'SUSPENDED',
    'TERMINATING', 'TERMINATED', 'ERROR',
)
kernelstatus = postgresql.ENUM(
    *kernelstatus_choices,
    name='kernelstatus')
def 升级():
    op.alter_column('kernels', column_name='status', 
                    type_=sa.Enum(*kernelstatus_choices, name='kernelstatus'),
                    postgresql_using='status::kernelstatus')

\n现在，我想将\'PENDING\'状态添加到kernelstatus类型中。所以我按照以下方式实现，参考了一些 文章。\n

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 升级():
    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;')

\n但它一直生成以下错误：\n

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

\n我已经尝试了向枚举类型添加值的解决方案，但这在与Alembic一起使用时不起作用，因为每个Alembic修订在一个事务中运行，而ALTER TYPE语句无法在事务中运行。此外，在PostgreSQL中没有删除枚举类型值的语句，因此在我的情况下，仅向枚举类型中添加一个值不可能是最终解决方案。\n有人能帮助我吗？

0

1 答案

0 Comments

问题出现的原因是因为在将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

0