在op.batch_alter_table中自动生成alter_colum

Question

我用@comparators.dispatch_for("table")编写了一个Alembic比较器来附加AlterColumnOp：

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

在迁移脚本中生成以下内容：

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

但是我需要像这样：batch_alter_table中获取它：

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

如何在comparator函数中实现此功能？

Answer 1

以下是替代解决方案：

@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}'))"
    )