Autogenerate alter_colum in op.batch_alter_table

Question

I wrote an alembic comparator with @comparators.dispatch_for("table") to append a AlterColumnOp:

modify_ops.ops.append(
    AlterColumnOp(
        table_name=tablename,
        column_name=column.name,
        modify_type=Enum(*current_column_enum_values),
    )
)

Which generate the following in migration script:

op.alter_column('user', 'status', type_=sa.Enum('ENABLED', 'DISABLED', 'NEW_ONE'))

But i need to get it in batch_alter_table, like this:

with op.batch_alter_table("user") as batch_op:
        batch_op.alter_column('status', type_=sa.Enum('ENABLED', 'DISABLED', 'NEW_ONE'))

How can i implement this in my comparator function ?

score 0 · Answer 1 · answered Feb 21 '20 at 10:53

This following is an alternative solution:

@Operations.register_operation("update_enum")
class UpdateEnumOp(MigrateOperation):
    def __init__(
        self, table_name: str, column_name: str, new_values: typing.List[str]
    ) -> None:
        self.table_name = table_name
        self.column_name = column_name
        self.new_values = new_values

    @classmethod
    def update_enum(
        cls, operations, table_name: str, column_name: str, new_values: typing.List[str]
    ) -> "UpdateEnumOp":
        op = UpdateEnumOp(table_name, column_name, new_values)
        return operations.invoke(op)

    def reverse(self):
        ...

@renderers.dispatch_for(UpdateEnumOp)
def render_update_enum(autogen_context, op: UpdateEnumOp) -> str:
    str_enum_values = "', '".join(op.new_values)
    return (
        f"with op.batch_alter_table(\"{op.table_name}\") as batch_op:\n"
        f"    batch_op.alter_column(\"{op.column_name}\", type_=sa.Enum('{str_enum_values}'))"
    )

Autogenerate alter_colum in op.batch_alter_table

1 Answers1