python - 使用 Alembic 更改枚举字段

Question

当使用早于 9.1 的 PostgreSQL 版本（为枚举添加 ALTER TYPE）时，如何在 alembic 迁移中将元素添加到 Enum 字段？这个SO question 解释了直接过程，但我不太确定如何最好地使用 alembic 翻译它。

这就是我所拥有的：

new_type = sa.Enum('nonexistent_executable', 'output_limit_exceeded',
                   'signal', 'success', 'timed_out', name='status')
old_type = sa.Enum('nonexistent_executable', 'signal', 'success', 'timed_out',
                   name='status')
tcr = sa.sql.table('testcaseresult',
                   sa.Column('status', new_type, nullable=False))


def upgrade():
    op.alter_column('testcaseresult', u'status', type_=new_type,
                    existing_type=old_type)


def downgrade():
    op.execute(tcr.update().where(tcr.c.status==u'output_limit_exceeded')
               .values(status='timed_out'))
    op.alter_column('testcaseresult', u'status', type_=old_type,
                    existing_type=new_type)

不幸的是，上述内容仅ALTER TABLE testcaseresult ALTER COLUMN status TYPE status在升级时产生，基本上什么都不做。

score 58 · Accepted Answer

我决定尝试尽可能直接地遵循postgres 方法，并提出了以下迁移。

from alembic import op
import sqlalchemy as sa

old_options = ('nonexistent_executable', 'signal', 'success', 'timed_out')
new_options = sorted(old_options + ('output_limit_exceeded',))

old_type = sa.Enum(*old_options, name='status')
new_type = sa.Enum(*new_options, name='status')
tmp_type = sa.Enum(*new_options, name='_status')

tcr = sa.sql.table('testcaseresult',
                   sa.Column('status', new_type, nullable=False))


def upgrade():
    # Create a tempoary "_status" type, convert and drop the "old" type
    tmp_type.create(op.get_bind(), checkfirst=False)
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE _status'
               ' USING status::text::_status')
    old_type.drop(op.get_bind(), checkfirst=False)
    # Create and convert to the "new" status type
    new_type.create(op.get_bind(), checkfirst=False)
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE status'
               ' USING status::text::status')
    tmp_type.drop(op.get_bind(), checkfirst=False)


def downgrade():
    # Convert 'output_limit_exceeded' status into 'timed_out'
    op.execute(tcr.update().where(tcr.c.status==u'output_limit_exceeded')
               .values(status='timed_out'))
    # Create a tempoary "_status" type, convert and drop the "new" type
    tmp_type.create(op.get_bind(), checkfirst=False)
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE _status'
               ' USING status::text::_status')
    new_type.drop(op.get_bind(), checkfirst=False)
    # Create and convert to the "old" status type
    old_type.create(op.get_bind(), checkfirst=False)
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE status'
               ' USING status::text::status')
    tmp_type.drop(op.get_bind(), checkfirst=False)

看来，alembicUSING在其alter_table方法中没有直接支持该声明。

score 26 · Accepted Answer

我使用了一种比我基于此的公认答案更简单的方法，步骤更少。在此示例中，我将假装有问题的枚举称为“status_enum”，因为在接受的答案中，列和枚举都使用“状态”使我感到困惑。

from alembic import op 
import sqlalchemy as sa

name = 'status_enum'
tmp_name = 'tmp_' + name

old_options = ('nonexistent_executable', 'signal', 'success', 'timed_out')
new_options = sorted(old_options + ('output_limit_exceeded',))

new_type = sa.Enum(*new_options, name=name)
old_type = sa.Enum(*old_options, name=name)

tcr = sa.sql.table('testcaseresult',
                   sa.Column('status', new_type, nullable=False))

def upgrade():
    op.execute('ALTER TYPE ' + name + ' RENAME TO ' + tmp_name)

    new_type.create(op.get_bind())
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status ' +
               'TYPE ' + name + ' USING status::text::' + name)
    op.execute('DROP TYPE ' + tmp_name)


def downgrade():
    # Convert 'output_limit_exceeded' status into 'timed_out'                                                                                                                      
    op.execute(tcr.update().where(tcr.c.status=='output_limit_exceeded')
               .values(status='timed_out'))

    op.execute('ALTER TYPE ' + name + ' RENAME TO ' + tmp_name)

    old_type.create(op.get_bind())
    op.execute('ALTER TABLE testcaseresult ALTER COLUMN status ' +
               'TYPE ' + name + ' USING status::text::' + name)
    op.execute('DROP TYPE ' + tmp_name)

score 18 · Accepted Answer

这运行没有问题：

from alembic import op

def upgrade():
    op.execute("COMMIT")
    op.execute("ALTER TYPE enum_type ADD VALUE 'new_value'")

def downgrade():
    ...

参考

score 13 · Accepted Answer

从 Postgres 9.1 开始，可以使用ALTER TYPE语句向枚举添加新值。由于不能在事务中完成，因此这很复杂。然而，这可以通过提交 alembic 的事务来解决，请参见此处。

实际上，我在使用旧的、更详细的解决方案时遇到了问题，因为 Postgres 无法自动转换列的默认值。

score 8 · Accepted Answer

我在尝试将列类型迁移到另一个时遇到了同样的问题。我使用以下要求：

Alembic==0.9.4
SQLAlchemy==1.1.12

您可以将参数提供postgresql_using为alembic.op.alter_column.

from alembic import op
import sqlalchemy as types

op.alter_column(
    table_name='my_table',
    column_name='my_column',
    type_=types.NewType,
    # allows to use postgresql USING
    postgresql_using="my_column::PostgesEquivalentOfNewType",
)

我希望它可以帮助。

score 4 · Accepted Answer

在直接 SQL 中，这适用于 Postgres，如果您的枚举中的事物的顺序不需要与上面完全相同：

ALTER TYPE status ADD value 'output_limit_exceeded' after 'timed_out';

score 4 · Accepted Answer

首先将列类型更改为 VARCHAR()。

然后删除您的类型并使用新字段创建新类型。

最后将您的列类型更改为新创建的类型。

def upgrade():
    op.execute(
        '''
        ALTER TABLE your_table ALTER COLUMN your_enum_column TYPE VARCHAR(255);

        DROP TYPE IF EXISTS your_enum_type;

        CREATE TYPE your_enum_type AS ENUM 
            ('value1', 'value2', 'value3', 'value4');

        ALTER TABLE your_table ALTER COLUMN your_enum_column TYPE your_enum_type 
            USING (your_enum_column::your_enum_type);
        '''
    )


def downgrade():
    op.execute(
        '''
        ALTER TABLE your_table ALTER COLUMN your_enum_column TYPE VARCHAR(255);

        DROP TYPE IF EXISTS your_enum_type;

        CREATE TYPE your_enum_type AS ENUM 
            ('value1', 'value2', 'value3');

        ALTER TABLE your_table ALTER COLUMN your_enum_column TYPE your_enum_type 
            USING (your_enum_column::your_enum_type);
        '''
    )

score 1 · Accepted Answer

我需要在迁移类型时移动数据，包括删除一些旧类型，所以我想我会根据（真棒）接受的答案（https://stackoverflow.com/a/14845740 ）写一个更通用的方法/629272）。希望这可以帮助同一条船上的其他人！

# This migration will move data from one column to two others based on the type
# for a given row, and modify the type of each row.
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision = '000000000001'
down_revision = '000000000000'
branch_labels = None
depends_on = None

# This set of options makes up the old type.
example_types_old = (
    'EXAMPLE_A',
    'EXAMPLE_B',
    'EXAMPLE_C',
)
example_type_enum_old = postgresql.ENUM(*example_types_old, name='exampletype')

# This set of options makes up the new type.
example_types_new = (
    'EXAMPLE_C',
    'EXAMPLE_D',
    'EXAMPLE_E',
)
example_type_enum_new = postgresql.ENUM(*example_types_new, name='exampletype')

# This set of options includes everything from the old and new types.
example_types_tmp = set(example_types_old + example_types_new)
example_type_enum_tmp = postgresql.ENUM(*example_types_tmp, name='_exampletype')

# This is a table view from which we can select and update as necessary. This
# only needs to include the relevant columns which are in either the old or new
# version of the table.
examples_view = sa.Table(
    # Use the name of the actual table so it is modified in the upgrade and
    # downgrade.
    'examples',
    sa.MetaData(),
    sa.Column('id', sa.Integer, primary_key=True),
    # Use the _tmp type so all types are usable.
    sa.Column('example_type', example_type_enum_tmp),
    # This is a column from which the data will be migrated, after which the
    # column will be removed.
    sa.Column('example_old_column', sa.Integer),
    # This is a column to which data from the old column will be added if the
    # type is EXAMPLE_A.
    sa.Column('example_new_column_a', sa.Integer),
    # This is a column to which data from the old column will be added if the
    # type is EXAMPLE_B.
    sa.Column('example_new_column_b', sa.Integer),
)


def upgrade():
    connection = op.get_bind()

    # Add the new column to which data will be migrated.
    example_new_column_a = sa.Column(
        'example_new_column_a',
        sa.Integer,
        nullable=True
    )
    op.add_column('examples', example_new_column_a)

    # Add the new column to which data will be migrated.
    example_new_column_b = sa.Column(
        'example_new_column_b',
        sa.Integer,
        nullable=True
    )
    op.add_column('examples', example_new_column_b)

    # Create the temporary enum and change the example_type column to use the
    # temporary enum.
    # The USING statement automatically maps the old enum to the temporary one.
    example_type_enum_tmp.create(connection, checkfirst=False)
    # Change to the temporary type and map from the old type to the temporary
    # one.
    op.execute('''
        ALTER TABLE examples
            ALTER COLUMN example_type
                TYPE _exampletype
                USING example_type::text::_exampletype
    ''')

    # Move data from example_old_column to example_new_column_a and change its
    # type to EXAMPLE_D if the type is EXAMPLE_A.
    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_A'
        ).values(
            example_type='EXAMPLE_D',
            example_new_column_a=examples_view.c.example_old_column,
        )
    )

    # Move data from example_old_column to example_new_column_b and change its
    # type to EXAMPLE_E if the type is EXAMPLE_B.
    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_B'
        ).values(
            example_type='EXAMPLE_E',
            example_new_column_b=examples_view.c.example_old_column,
        )
    )

    # Move any remaining data from example_old_column to example_new_column_a
    # and keep its type as EXAMPLE_C.
    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_C'
        ).values(
            example_type='EXAMPLE_C',
            example_new_column_a=examples_view.c.example_old_column,
        )
    )

    # Delete the old enum now that the data with the old types have been moved.
    example_type_enum_old.drop(connection, checkfirst=False)

    # Create the new enum and change the example_type column to use the new
    # enum.
    # The USING statement automatically maps the temporary enum to the new one.
    example_type_enum_new.create(connection, checkfirst=False)
    op.execute('''
        ALTER TABLE examples
            ALTER COLUMN example_type
                TYPE exampletype
                USING example_type::text::exampletype
    ''')

    # Delete the temporary enum.
    example_type_enum_tmp.drop(connection, checkfirst=False)

    # Remove the old column.
    op.drop_column('examples', 'example_old_column')


# The downgrade just performs the opposite of all the upgrade operations but in
# reverse.
def downgrade():
    connection = op.get_bind()

    example_old_column = sa.Column(
        'example_old_column',
        sa.Integer,
        nullable=True
    )
    op.add_column('examples', example_old_column)

    example_type_enum_tmp.create(connection, checkfirst=False)
    op.execute('''
        ALTER TABLE examples
            ALTER COLUMN example_type
                TYPE _exampletype
                USING example_type::text::_exampletype
    ''')

    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_C'
        ).values(
            example_type='EXAMPLE_C',
            example_old_column=examples_view.c.example_new_column_b,
        )
    )

    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_E'
        ).values(
            example_type='EXAMPLE_B',
            example_old_column=examples_view.c.example_new_column_b,
        )
    )

    connection.execute(
        examples_view.update().where(
            examples_view.c.example_type == 'EXAMPLE_D'
        ).values(
            example_type='EXAMPLE_A',
            example_old_column=examples_view.c.example_new_column_a,
        )
    )

    example_type_enum_old.create(connection, checkfirst=False)
    op.execute('''
        ALTER TABLE examples
            ALTER COLUMN example_type
                TYPE exampletype
                USING example_type::text::exampletype
    ''')

    example_type_enum_tmp.drop(connection, checkfirst=False)

    op.drop_column('examples', 'example_new_column_b')
    op.drop_column('examples', 'example_new_column_a')

score 1 · Accepted Answer

找到另一个方便的方法

op.execute('ALTER TYPE enum_type ADD VALUE new_value')
op.execute('ALTER TYPE enum_type ADD VALUE new_value BEFORE old_value')
op.execute('ALTER TYPE enum_type ADD VALUE new_value AFTER old_value')

score 1 · Accepted Answer

此方法可用于更新枚举

def upgrade():
    op.execute("ALTER TYPE categorytype RENAME VALUE 'EXAMPLE_A' TO 'EXAMPLE_B'")


def downgrade():
    op.execute("ALTER TYPE categorytype RENAME VALUE 'EXAMPLE_B' TO 'EXAMPLE_A'")

score 1 · Accepted Answer

这种方法类似于公认的解决方案，但有细微差别：

它使用op.batch_alter_table而不是op.execute('ALTER TABLE')，因此此解决方案适用于 PostgreSQL 和 SQLite：SQLite 不支持ALTER TABLE，但 alembic 提供了对它的支持op.batch_alter_table
它不使用原始 SQL

from alembic import op
import sqlalchemy as sa


# Describing of enum
enum_name = "status"
temp_enum_name = f"temp_{enum_name}"
old_values = ("nonexistent_executable", "signal", "success", "timed_out")
new_values = ("output_limit_exceeded", *old_values)
downgrade_to = ("output_limit_exceeded", "timed_out") # on downgrade convert [0] to [1]
old_type = sa.Enum(*old_values, name=enum_name)
new_type = sa.Enum(*new_values, name=enum_name)
temp_type = sa.Enum(*new_values, name=temp_enum_name)

# Describing of table
table_name = "testcaseresult"
column_name = "status"
temp_table = sa.sql.table(
    table_name,
    sa.Column(
        column_name,
        new_type,
        nullable=False
    )
)


def upgrade():
    # temp type to use instead of old one
    temp_type.create(op.get_bind(), checkfirst=False)

    # changing of column type from old enum to new one.
    # SQLite will create temp table for this
    with op.batch_alter_table(table_name) as batch_op:
        batch_op.alter_column(
            column_name,
            existing_type=old_type,
            type_=temp_type,
            existing_nullable=False,
            postgresql_using=f"{column_name}::text::{temp_enum_name}"
        )

    # remove old enum, create new enum
    old_type.drop(op.get_bind(), checkfirst=False)
    new_type.create(op.get_bind(), checkfirst=False)

    # changing of column type from temp enum to new one.
    # SQLite will create temp table for this
    with op.batch_alter_table(table_name) as batch_op:
        batch_op.alter_column(
            column_name,
            existing_type=temp_type,
            type_=new_type,
            existing_nullable=False,
            postgresql_using=f"{column_name}::text::{enum_name}"
        )

    # remove temp enum
    temp_type.drop(op.get_bind(), checkfirst=False)


def downgrade():
    # old enum don't have new value anymore.
    # before downgrading from new enum to old one,
    # we should replace new value from new enum with
    # somewhat of old values from old enum
    op.execute(
        temp_table
        .update()
        .where(
            temp_table.c.status == downgrade_to[0]
        )
        .values(
            status=downgrade_to[1]
        )
    )

    temp_type.create(op.get_bind(), checkfirst=False)

    with op.batch_alter_table(table_name) as batch_op:
        batch_op.alter_column(
            column_name,
            existing_type=new_type,
            type_=temp_type,
            existing_nullable=False,
            postgresql_using=f"{column_name}::text::{temp_enum_name}"
        )

    new_type.drop(op.get_bind(), checkfirst=False)
    old_type.create(op.get_bind(), checkfirst=False)

    with op.batch_alter_table(table_name) as batch_op:
        batch_op.alter_column(
            column_name,
            existing_type=temp_type,
            type_=old_type,
            existing_nullable=False,
            postgresql_using=f"{column_name}::text::{enum_name}"
        )

    temp_type.drop(op.get_bind(), checkfirst=False)

从公认的解决方案：

似乎 alembic 在其 alter_table 方法中没有直接支持 USING 语句。

目前，alembicUSING在其alter_table方法中支持声明。

score 0 · Accepted Answer

由于我遇到了转换错误和默认值问题，因此我根据接受的答案写了一个更通用的答案：

def replace_enum_values(
        name: str,
        old: [str],
        new: [str],
        modify: [(str, str, str)]
):
    """
    Replaces an enum's list of values.

    Args:
        name: Name of the enum
        new: New list of values
        old: Old list of values
        modify: List of tuples of table name
        and column to modify (which actively use the enum).
        Assumes each column has a default val.
    """
    connection = op.get_bind()

    tmp_name = "{}_tmp".format(name)

    # Rename old type
    op.execute(
        "ALTER TYPE {} RENAME TO {};"
        .format(name, tmp_name)
    )

    # Create new type
    lsl = sa.Enum(*new, name=name)
    lsl.create(connection)

    # Replace all usages
    for (table, column) in modify:
        # Get default to re-set later
        default_typed = connection.execute(
            "SELECT column_default "
            "FROM information_schema.columns "
            "WHERE table_name='{table}' "
            "AND column_name='{column}';"
            .format(table=table, column=column)
        ).first()[0]  # type: str

        # Is bracketed already
        default = default_typed[:default_typed.index("::")]

        # Set all now invalid values to default
        connection.execute(
            "UPDATE {table} "
            "SET {column}={default} "
            "WHERE {column} NOT IN {allowed};"
            .format(
                table=table,
                column=column,
                # Invalid: What isn't contained in both new and old
                # Can't just remove what's not in new because we get
                # a type error
                allowed=tuple(set(old).intersection(set(new))),
                default=default
            )
        )

        op.execute(
            "ALTER TABLE {table} "
            # Default needs to be dropped first
            "ALTER COLUMN {column} DROP DEFAULT,"
            # Replace the tpye
            "ALTER COLUMN {column} TYPE {enum_name} USING {column}::text::{enum_name},"
            # Reset default
            "ALTER COLUMN {column} SET DEFAULT {default};"
            .format(
                table=table,
                column=column,
                enum_name=name,
                default=default
            )
        )

    # Remove old type
    op.execute("DROP TYPE {};".format(tmp_name))

这可以从升级/降级中调用，如下所示：

replace_enum_values(
    name='enum_name',
    new=["A", "B"],
    old=["A", "C"],
    modify=[('some_table', 'some_column')]
)

所有无效值都将设置为 server_default。

score 0 · Accepted Answer

观察

为了减轻迁移时的痛苦，即使使用 PostgreSQL，我也总是使用非本地枚举。

非原生枚举只是带有约束的字符串，如果你编辑一个枚举，只有三种情况：

重命名枚举值
删除枚举值
添加枚举值。

对于迁移，2 和 3 是一对。这是可以理解的：如果你升级是为了添加，那么你必须在降级时删除，反之亦然。因此，让我们将它们分为两种类型。

执行

如果要重命名，通常我会将其分为三个步骤：

放弃旧的约束
将行的旧值更新为新值
创建新约束

在 alembic 中，这是通过以下方式完成的：

def update_enum(
    table, column, enum_class_name, target_values, olds_to_remove, news_to_add
):
    op.drop_constraint(f"ck_{table}_{enum_class_name}", table)

    for sql in update_enum_sqls(table, column, olds_to_remove, news_to_add):
        op.execute(sql)

    op.create_check_constraint(
        enum_class_name, table, sa.sql.column(column).in_(target_values)
    )

让我们先忘记update_enum_sqls它，将其用作 SQL 生成器。

如果它正在删除，那么仍然有三个步骤：

放弃旧的约束
删除具有旧值的行
创建新约束

所以基本上只有update_enum_sqls可能的行为不同。

如果是添加，只需两步：

放弃旧的约束
创建新约束

不过，我们可以忽略update_enum_sqls.

那么如何实现呢？没那么难...

def update_enum_sql(table, column, old_value, new_value):
    if new_value is not None:
        return f"UPDATE {table} SET {column} = '{new_value}' where {column} = '{old_value}'"
    else:
        return f"DELETE FROM {table} where {column} = '{old_value}'"


def update_enum_sqls(table, column, olds_to_remove, news_to_add):
    if len(olds_to_remove) != len(news_to_add):
        raise NotImplementedError
    return [
        update_enum_sql(table, column, old, new)
        for old, new in zip(olds_to_remove, news_to_add)
    ]

例子

既然我们准备了配料，让我们申请：

def upgrade():
    # rename enum
    update_enum(
        "my_table",
        "my_enum",
        "myenumclassname",
        ["NEW", "ENUM", "VALUES"],
        ["OLD"],
        ["NEW"],
    )

    # add enum
    update_enum(
        "my_table",
        "my_enum",
        "myenumclassname",
        ["NEW", "ENUM", "VALUES"],
        [],
        [],
    )


def downgrade():
    # remove enum
    update_enum(
        "my_table",
        "my_enum",
        "myenumclassname",
        ["ENUM", "VALUES"],
        ["NEW"],
        [None],  # this will delete rows with "NEW", USE WITH CARE!!!
    )

    # edit enum
    update_enum(
        "my_table",
        "my_enum",
        "myenumclassname",
        ["OLD", "ENUM", "VALUES"],
        ["NEW"],
        ["OLD"],
    )

上面的代码也可以在gist上找到。

score -1 · Accepted Answer

该解决方案易于理解，并且对于升级和降级都非常有效。我已经以更详细的方式写了这个答案。

假设我们的enum_type外观是这样的：

enum_type = ('some_value_1', 'some_value_2')

我想enum_type通过添加一个新的枚举来改变，使它变成这样：

enum_type = ('some_value_1', 'some_value_2', 'new_value')

这可以通过以下方式完成：

from alembic import op


def upgrade():
    op.execute("COMMIT")
    op.execute("ALTER TYPE enum_type ADD VALUE 'new_value'")


def downgrade():
    # Drop 'new_value' from enum_type
    op.execute("ALTER TYPE enum_type RENAME TO enum_type_tmp")

    op.execute("CREATE TYPE enum_type AS ENUM('some_value_1', 'some_value_1')")

    op.execute("DROP TYPE enum_type_tmp")

注意：在降级过程中，如果您enum_type在表格中使用，那么您可以修改降级方法，如下所述：

def downgrade():
    # Drop 'new_value' from enum_type
    op.execute("UPDATE table_name"
               " SET column_name_using_enum_type_value = NULL"
               " WHERE column_name_using_enum_type_value = 'new_value'")    

    op.execute("ALTER TYPE enum_type RENAME TO enum_type_tmp")

    op.execute("CREATE TYPE enum_type AS ENUM('some_value_1', 'some_value_1')")

    op.execute("ALTER TABLE table_name"
               " ALTER COLUMN column_name_using_enum_type_value TYPE enum_type"
               " USING column_name_using_enum_type_value::text::enum_type")

    op.execute("DROP TYPE enum_type_tmp")

python - 使用 Alembic 更改枚举字段

14 回答 14

观察

执行

例子

Related

Reference