Автогенерация alter_colum в op.batch_alter_table - PullRequest
Новая
       46

Автогенерация alter_colum в op.batch_alter_table

0 голосов
/ 21 февраля 2020

Я написал компилятор alembi c с @comparators.dispatch_for("table") для добавления 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?

1 Ответ

0 голосов
/ 21 февраля 2020

Это следующее альтернативное решение: 

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