Demystifying Alembic's Handling of Enum Types in Table Alterations

When working with Alembic, the database migration tool for Python, there is an important aspect to consider regarding Enum types. Alembic does not automatically create an enum type when altering a table. Instead, you need to handle the creation manually.

Creating the Enum

To create an enum type in Postgres using Alembic, we need to define the enum and instruct the database to create it if it doesn't already exist. Let's take a look at the code snippet below:

# migrations/versions/add_role_to_user.py
def upgrade() -> None:
    with op.batch_alter_table("user", schema=None) as batch_op:
        permission_enum = sa.Enum(
            "USER", "ADMIN", "SUPERUSER", name="permission"
        )
        permission_enum.create(batch_op.get_bind(), checkfirst=True)
        batch_op.add_column(
            sa.Column(
                "role", permission_enum, server_default="USER", nullable=False
            )
        )

In the upgrade() function, we create an enum called permission with values "USER", "ADMIN", and "SUPERUSER". The name parameter specifies the name of the enum. We then create the enum using permission_enum.create() and add a new column named "role" to the "user" table, specifying the permission_enum as its type.

The checkfirst=True parameter ensures that the permission_enum is only created if it doesn't already exist in the database.

To know more about the batch_alter_table, check out my article on configuring SQLModel and Alembic.

Handling Downgrades

In the downgrade() function, we need to handle the removal of the enum from the database when rolling back the migration.

# migrations/versions/add_role_to_user.py
def downgrade() -> None:
    with op.batch_alter_table("user", schema=None) as batch_op:
        batch_op.drop_column("role")
        permission_enum = sa.Enum(
            "USER", "ADMIN", "SUPERUSER", name="permission"
        )
        permission_enum.drop(batch_op.get_bind(), checkfirst=True)

In the downgrade() function, we use the batch_alter_table context manager to define the changes needed for the rollback. Here, we drop the "role" column added in the upgrade() function. Then, we define the permission_enum enum again and call permission_enum.drop() to remove the enum from the database.

As before, the checkfirst=True parameter ensures that the permission_enum is only dropped if it exists in the database.

Considerations and Alternatives

While the usage of enums simplifies the representation of discrete values, it's worth noting that alternative approaches such as integer or bit field flags may provide better flexibility and performance in some scenarios.