diff --git a/alembic_postgresql_enum/add_postgres_using_to_text.py b/alembic_postgresql_enum/add_postgres_using_to_text.py index adf33ba..8ac6593 100644 --- a/alembic_postgresql_enum/add_postgres_using_to_text.py +++ b/alembic_postgresql_enum/add_postgres_using_to_text.py @@ -38,18 +38,30 @@ def _postgres_using_alter_column(autogen_context: AutogenContext, op: ops.AlterC log = logging.getLogger(f"alembic.{__name__}") -def add_postgres_using_to_alter_operation(op: AlterColumnOp): +def add_postgres_using_to_alter_operation_text_to_enum(op: AlterColumnOp): assert op.modify_type is not None op.kw["postgresql_using"] = f"{op.column_name}::{op.modify_type.name}" log.info("postgresql_using added to %r.%r alteration", op.table_name, op.column_name) op.__class__ = PostgresUsingAlterColumnOp +def add_postgres_using_to_alter_operation_enum_to_enum(op: AlterColumnOp): + assert op.modify_type is not None + op.kw["postgresql_using"] = f"{op.column_name}::text::{op.modify_type.name}" + log.info("postgresql_using added to %r.%r alteration", op.table_name, op.column_name) + op.__class__ = PostgresUsingAlterColumnOp + + def add_postgres_using_to_text(upgrade_ops: UpgradeOps): """Add postgresql_using to alter_column expressions that changes type from string to enum""" for group_op in upgrade_ops.ops: - if isinstance(group_op, ModifyTableOps): - for i, op in enumerate(group_op.ops): - if isinstance(op, AlterColumnOp): - if isinstance(op.existing_type, String) and column_type_is_enum(op.modify_type): - add_postgres_using_to_alter_operation(op) + if not isinstance(group_op, ModifyTableOps): + continue + for i, op in enumerate(group_op.ops): + if not isinstance(op, AlterColumnOp): + continue + if column_type_is_enum(op.modify_type): + if column_type_is_enum(op.existing_type): + add_postgres_using_to_alter_operation_enum_to_enum(op) + elif isinstance(op.existing_type, String): + add_postgres_using_to_alter_operation_text_to_enum(op) diff --git a/tests/test_alter_column/test_rename_enum.py b/tests/test_alter_column/test_rename_enum.py new file mode 100644 index 0000000..d32840d --- /dev/null +++ b/tests/test_alter_column/test_rename_enum.py @@ -0,0 +1,79 @@ +from enum import Enum +from typing import TYPE_CHECKING + +from sqlalchemy import MetaData, Table, Column, insert +from sqlalchemy.dialects import postgresql + +from tests.base.run_migration_test_abc import CompareAndRunTestCase + +if TYPE_CHECKING: + from sqlalchemy import Connection + + +class OldEnum(Enum): + A = "a" + B = "b" + C = "c" + + +class NewEnum(Enum): + A = "a" + B = "b" + C = "c" + + +class TestRenameEnum(CompareAndRunTestCase): + """ + #115 + Unfortunately there is no way to add postgresql_using to alter_column in downgrade + """ + + disable_running = True + + def get_database_schema(self) -> MetaData: + database_schema = MetaData() + Table("a", database_schema, Column("value", postgresql.ENUM(OldEnum))) + return database_schema + + def get_target_schema(self) -> MetaData: + target_schema = MetaData() + Table("a", target_schema, Column("value", postgresql.ENUM(NewEnum))) + return target_schema + + def insert_migration_data(self, connection: "Connection", database_schema: MetaData) -> None: + a_table = database_schema.tables["a"] + connection.execute( + insert(a_table).values( + [ + {"value": OldEnum.A.name}, + {"value": OldEnum.B.name}, + {"value": OldEnum.B.name}, + {"value": OldEnum.C.name}, + ] + ) + ) + + def get_expected_upgrade(self) -> str: + return """ + # ### commands auto generated by Alembic - please adjust! ### + sa.Enum('A', 'B', 'C', name='newenum').create(op.get_bind()) + op.alter_column('a', 'value', + existing_type=postgresql.ENUM('A', 'B', 'C', name='oldenum'), + type_=postgresql.ENUM('A', 'B', 'C', name='newenum'), + existing_nullable=True, + postgresql_using='value::text::newenum') + sa.Enum('A', 'B', 'C', name='oldenum').drop(op.get_bind()) + # ### end Alembic commands ### + """ + + def get_expected_downgrade(self) -> str: + return """ + # ### commands auto generated by Alembic - please adjust! ### + sa.Enum('A', 'B', 'C', name='oldenum').create(op.get_bind()) + op.alter_column('a', 'value', + existing_type=postgresql.ENUM('A', 'B', 'C', name='newenum'), + type_=postgresql.ENUM('A', 'B', 'C', name='oldenum'), + existing_nullable=True) + sa.Enum('A', 'B', 'C', name='newenum').drop(op.get_bind()) + # ### end Alembic commands ### + """