Updating enum values in non-default schema fails

alembic_postgresql_enum/enum_alteration.py

@@ -206,7 +206,7 @@ def _set_enum_values(cls,
 
         drop_comparison_operators(connection, schema, enum_name, temporary_enum_name)
         connection.execute(sqlalchemy.text(
-            f"""DROP TYPE {temporary_enum_name};"""
+            f"""DROP TYPE {schema}.{temporary_enum_name};"""
         ))
 
     @classmethod

tests/test_enum_alteration.py → tests/sync_enum_values/test_array_column.py

@@ -10,11 +10,7 @@
 if TYPE_CHECKING:
     from sqlalchemy import Connection
 
-from tests.schemas import (get_schema_with_enum_variants,
-                           get_schema_with_enum_in_array_variants,
-                           USER_TABLE_NAME,
-                           USER_STATUS_ENUM_NAME,
-                           USER_STATUS_COLUMN_NAME,
+from tests.schemas import (get_schema_with_enum_in_array_variants,
                            DEFAULT_SCHEMA,
                            CAR_TABLE_NAME,
                            CAR_COLORS_COLUMN_NAME,
@@ -23,135 +19,6 @@
 from tests.utils.migration_context import create_migration_context
 
 
-def test_add_new_enum_value_render(connection: 'Connection'):
-    """Check that enum variants are updated when new variant is added"""
-    old_enum_variants = ["active", "passive"]
-
-    database_schema = get_schema_with_enum_variants(old_enum_variants)
-    database_schema.create_all(connection)
-
-    new_enum_variants = old_enum_variants.copy()
-    new_enum_variants.append('banned')
-
-    target_schema = get_schema_with_enum_variants(new_enum_variants)
-
-    context = create_migration_context(connection, target_schema)
-
-    template_args = {}
-    autogenerate._render_migration_diffs(context, template_args)
-
-    assert (template_args["upgrades"] ==
-            f"""# ### commands auto generated by Alembic - please adjust! ###
-    op.sync_enum_values('{DEFAULT_SCHEMA}', '{USER_STATUS_ENUM_NAME}', [{', '.join(map(repr, new_enum_variants))}],
-                        [('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')],
-                        enum_values_to_rename=[])
-    # ### end Alembic commands ###""")
-    assert (template_args["downgrades"] ==
-            f"""# ### commands auto generated by Alembic - please adjust! ###
-    op.sync_enum_values('{DEFAULT_SCHEMA}', '{USER_STATUS_ENUM_NAME}', [{', '.join(map(repr, old_enum_variants))}],
-                        [('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')],
-                        enum_values_to_rename=[])
-    # ### end Alembic commands ###""")
-
-
-def test_add_new_enum_value_diff_tuple(connection: 'Connection'):
-    """Check that enum variants are updated when new variant is added"""
-    old_enum_variants = ["active", "passive"]
-
-    database_schema = get_schema_with_enum_variants(old_enum_variants)
-    database_schema.create_all(connection)
-
-    new_enum_variants = old_enum_variants.copy()
-    new_enum_variants.append('banned')
-
-    target_schema = get_schema_with_enum_variants(new_enum_variants)
-
-    context = create_migration_context(connection, target_schema)
-
-    autogen_context = api.AutogenContext(context, target_schema)
-
-    uo = ops.UpgradeOps(ops=[])
-    autogenerate._produce_net_changes(autogen_context, uo)
-
-    diffs = uo.as_diffs()
-    assert len(diffs) == 1
-    sync_diff_tuple = diffs[0]
-
-    assert sync_diff_tuple == (
-        SyncEnumValuesOp.operation_name,
-        old_enum_variants,
-        new_enum_variants,
-        [(USER_TABLE_NAME, USER_STATUS_COLUMN_NAME)]
-    )
-
-
-def test_remove_enum_value_diff_tuple(connection: 'Connection'):
-    """Check that enum variants are updated when new variant is removed"""
-    old_enum_variants = ["active", "passive", "banned"]
-
-    database_schema = get_schema_with_enum_variants(old_enum_variants)
-    database_schema.create_all(connection)
-
-    new_enum_variants = old_enum_variants.copy()
-    new_enum_variants.remove('banned')
-
-    target_schema = get_schema_with_enum_variants(new_enum_variants)
-
-    context = create_migration_context(connection, target_schema)
-
-    autogen_context = api.AutogenContext(context, target_schema)
-
-    uo = ops.UpgradeOps(ops=[])
-    autogenerate._produce_net_changes(autogen_context, uo)
-
-    diffs = uo.as_diffs()
-    assert len(diffs) == 1
-
-    change_variants_diff_tuple = diffs[0]
-    operation_name, old_values, new_values, affected_columns = change_variants_diff_tuple
-
-    assert operation_name == SyncEnumValuesOp.operation_name
-    assert old_values == old_enum_variants
-    assert new_values == new_enum_variants
-    assert affected_columns == [
-        (USER_TABLE_NAME, USER_STATUS_COLUMN_NAME)
-    ]
-
-
-def test_rename_enum_value_diff_tuple(connection: 'Connection'):
-    """Check that enum variants are updated when a variant is renamed"""
-    old_enum_variants = ["active", "passive", "banned"]
-
-    database_schema = get_schema_with_enum_variants(old_enum_variants)
-    database_schema.create_all(connection)
-
-    new_enum_variants = old_enum_variants.copy()
-    new_enum_variants.remove('banned')
-    new_enum_variants.append('inactive')
-
-    target_schema = get_schema_with_enum_variants(new_enum_variants)
-
-    context = create_migration_context(connection, target_schema)
-
-    autogen_context = api.AutogenContext(context, target_schema)
-
-    uo = ops.UpgradeOps(ops=[])
-    autogenerate._produce_net_changes(autogen_context, uo)
-
-    diffs = uo.as_diffs()
-    assert len(diffs) == 1
-
-    change_variants_diff_tuple = diffs[0]
-    operation_name, old_values, new_values, affected_columns = change_variants_diff_tuple
-
-    assert operation_name == SyncEnumValuesOp.operation_name
-    assert old_values == old_enum_variants
-    assert new_values == new_enum_variants
-    assert affected_columns == [
-        (USER_TABLE_NAME, USER_STATUS_COLUMN_NAME)
-    ]
-
-
 def test_add_new_enum_value_render_with_array(connection: 'Connection'):
     """Check that enum variants are updated when new variant is added"""
     old_enum_variants = ['black', 'white', 'red', 'green', 'blue', 'other']

tests/sync_enum_values/test_render.py

@@ -0,0 +1,149 @@
+from typing import TYPE_CHECKING
+
+from alembic import autogenerate
+from alembic.autogenerate import api
+from alembic.operations import ops
+
+from alembic_postgresql_enum.enum_alteration import SyncEnumValuesOp
+
+if TYPE_CHECKING:
+    from sqlalchemy import Connection
+
+from tests.schemas import (get_schema_with_enum_variants,
+                           USER_TABLE_NAME,
+                           USER_STATUS_ENUM_NAME,
+                           USER_STATUS_COLUMN_NAME,
+                           DEFAULT_SCHEMA
+                           )
+from tests.utils.migration_context import create_migration_context
+
+
+def test_add_new_enum_value_render(connection: 'Connection'):
+    """Check that enum variants are updated when new variant is added"""
+    old_enum_variants = ["active", "passive"]
+
+    database_schema = get_schema_with_enum_variants(old_enum_variants)
+    database_schema.create_all(connection)
+
+    new_enum_variants = old_enum_variants.copy()
+    new_enum_variants.append('banned')
+
+    target_schema = get_schema_with_enum_variants(new_enum_variants)
+
+    context = create_migration_context(connection, target_schema)
+
+    template_args = {}
+    autogenerate._render_migration_diffs(context, template_args)
+
+    assert (template_args["upgrades"] ==
+            f"""# ### commands auto generated by Alembic - please adjust! ###
+    op.sync_enum_values('{DEFAULT_SCHEMA}', '{USER_STATUS_ENUM_NAME}', [{', '.join(map(repr, new_enum_variants))}],
+                        [('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')],
+                        enum_values_to_rename=[])
+    # ### end Alembic commands ###""")
+    assert (template_args["downgrades"] ==
+            f"""# ### commands auto generated by Alembic - please adjust! ###
+    op.sync_enum_values('{DEFAULT_SCHEMA}', '{USER_STATUS_ENUM_NAME}', [{', '.join(map(repr, old_enum_variants))}],
+                        [('{USER_TABLE_NAME}', '{USER_STATUS_COLUMN_NAME}')],
+                        enum_values_to_rename=[])
+    # ### end Alembic commands ###""")
+
+
+def test_add_new_enum_value_diff_tuple(connection: 'Connection'):
+    """Check that enum variants are updated when new variant is added"""
+    old_enum_variants = ["active", "passive"]
+
+    database_schema = get_schema_with_enum_variants(old_enum_variants)
+    database_schema.create_all(connection)
+
+    new_enum_variants = old_enum_variants.copy()
+    new_enum_variants.append('banned')
+
+    target_schema = get_schema_with_enum_variants(new_enum_variants)
+
+    context = create_migration_context(connection, target_schema)
+
+    autogen_context = api.AutogenContext(context, target_schema)
+
+    uo = ops.UpgradeOps(ops=[])
+    autogenerate._produce_net_changes(autogen_context, uo)
+
+    diffs = uo.as_diffs()
+    assert len(diffs) == 1
+    sync_diff_tuple = diffs[0]
+
+    assert sync_diff_tuple == (
+        SyncEnumValuesOp.operation_name,
+        old_enum_variants,
+        new_enum_variants,
+        [(USER_TABLE_NAME, USER_STATUS_COLUMN_NAME)]
+    )
+
+
+def test_remove_enum_value_diff_tuple(connection: 'Connection'):
+    """Check that enum variants are updated when new variant is removed"""
+    old_enum_variants = ["active", "passive", "banned"]
+
+    database_schema = get_schema_with_enum_variants(old_enum_variants)
+    database_schema.create_all(connection)
+
+    new_enum_variants = old_enum_variants.copy()
+    new_enum_variants.remove('banned')
+
+    target_schema = get_schema_with_enum_variants(new_enum_variants)
+
+    context = create_migration_context(connection, target_schema)
+
+    autogen_context = api.AutogenContext(context, target_schema)
+
+    uo = ops.UpgradeOps(ops=[])
+    autogenerate._produce_net_changes(autogen_context, uo)
+
+    diffs = uo.as_diffs()
+    assert len(diffs) == 1
+
+    change_variants_diff_tuple = diffs[0]
+    operation_name, old_values, new_values, affected_columns = change_variants_diff_tuple
+
+    assert operation_name == SyncEnumValuesOp.operation_name
+    assert old_values == old_enum_variants
+    assert new_values == new_enum_variants
+    assert affected_columns == [
+        (USER_TABLE_NAME, USER_STATUS_COLUMN_NAME)
+    ]
+
+
+def test_rename_enum_value_diff_tuple(connection: 'Connection'):
+    """Check that enum variants are updated when a variant is renamed"""
+    old_enum_variants = ["active", "passive", "banned"]
+
+    database_schema = get_schema_with_enum_variants(old_enum_variants)
+    database_schema.create_all(connection)
+
+    new_enum_variants = old_enum_variants.copy()
+    new_enum_variants.remove('banned')
+    new_enum_variants.append('inactive')
+
+    target_schema = get_schema_with_enum_variants(new_enum_variants)
+
+    context = create_migration_context(connection, target_schema)
+
+    autogen_context = api.AutogenContext(context, target_schema)
+
+    uo = ops.UpgradeOps(ops=[])
+    autogenerate._produce_net_changes(autogen_context, uo)
+
+    diffs = uo.as_diffs()
+    assert len(diffs) == 1
+
+    change_variants_diff_tuple = diffs[0]
+    operation_name, old_values, new_values, affected_columns = change_variants_diff_tuple
+
+    assert operation_name == SyncEnumValuesOp.operation_name
+    assert old_values == old_enum_variants
+    assert new_values == new_enum_variants
+    assert affected_columns == [
+        (USER_TABLE_NAME, USER_STATUS_COLUMN_NAME)
+    ]
+
+

tests/sync_enum_values/test_run_in_different_schema.py

@@ -0,0 +1,59 @@
+# Case No. 2 from https://github.com/Pogchamp-company/alembic-postgresql-enum/issues/26
+import enum
+from typing import TYPE_CHECKING
+
+from alembic.operations import Operations
+from alembic.runtime.migration import MigrationContext
+
+from alembic_postgresql_enum import get_defined_enums
+from tests.schemas import ANOTHER_SCHEMA_NAME
+
+if TYPE_CHECKING:
+    from sqlalchemy import Connection
+from sqlalchemy import MetaData, Column, Integer
+from sqlalchemy.orm import declarative_base
+from sqlalchemy.dialects.postgresql import ENUM
+
+my_metadata = MetaData(schema=ANOTHER_SCHEMA_NAME)
+
+
+Base = declarative_base(metadata=my_metadata)
+
+
+class _TestStatus(enum.Enum):
+    PENDING = "PENDING"
+    SUCCESS = "SUCCESS"
+    FAILED = "FAILED"
+
+
+class TableWithExplicitEnumSchema(Base):
+    __tablename__ = "test"
+    id = Column(Integer, primary_key=True)
+
+    status = Column(
+        ENUM(_TestStatus, name="test_status", schema=Base.metadata.schema),
+        nullable=False,
+    )
+
+
+def test_run_in_different_schema(connection: 'Connection'):
+    """https://github.com/Pogchamp-company/alembic-postgresql-enum/issues/34"""
+    old_enum_variants = list(map(lambda item: item.name, _TestStatus))
+
+    database_schema = my_metadata
+    database_schema.create_all(connection)
+
+    new_enum_variants = old_enum_variants.copy()
+    new_enum_variants.append('WAITING_FOR_APPROVAL')
+
+    mc = MigrationContext.configure(connection)
+    ops = Operations(mc)
+
+    ops.sync_enum_values(Base.metadata.schema, 'test_status', new_enum_variants,
+                         [(TableWithExplicitEnumSchema.__tablename__, 'status')])
+
+    defined = get_defined_enums(connection, Base.metadata.schema)
+
+    assert defined == {
+        'test_status': tuple(new_enum_variants)
+    }