Alembic issues with postgresql enum processing

Alexandrov Roman — Sun, 06 Aug 2023 07:42:20 +0000

Disclaimer

In this article, I would like to focus specifically on the problems of alembic, and not on the appropriateness of using enumerations in general. I will also not go into detail on how to use them.

A brief preface

Let us have an enum order_status_enum:

from enum import Enum

class OrderStatusEnum(Enum):
    WAITING_FOR_WORKER = 'WAITING_FOR_WORKER'
    IN_PROGRESS = 'IN_PROGRESS'
    DONE = 'DONE'

And the next use of it:

from sqlalchemy import Integer
from sqlalchemy.dialects.postgresql import ENUM as PgEnum

from app.database import BaseModel


class Order(BaseModel):
    id = Column(Integer, primary_key=True, autoincrement=True)

    status = Column(PgEnum(OrderStatusEnum, name='order_status_enum'), nullable=False, default=OrderStatusEnum.WAITING_FOR_WORKER)
    # ... other fields

Migration issues

Creation

Now that we have an enum, we can finally talk about migration issues. If when creating an enum do not specify create_type=False parameter, it may seem that everything is fine and works as it should. The migration is created and applied:

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('order',
                    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
                    sa.Column('status',
                              postgresql.ENUM('WAITING_FOR_WORKER', 'IN_PROGRESS', 'DONE', name='order_status_enum'),
                              nullable=True),
                    sa.PrimaryKeyConstraint('id', name=op.f('order_pkey'))
                    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('order')
    # ### end Alembic commands ###

But that's not really the case. If we downgrade and then try to upgrade again, it turns out that the downgrade didn't remove the enum. So how can we fix this? We need to add a line that will remove our enum. It will look like this as follows (for the sake of code consistency, I also wrote an explicit creation):

def _get_order_status():
    return postgresql.ENUM('WAITING_FOR_WORKER', 'IN_PROGRESS',
                           'DONE', name='order_status_enum')


def upgrade():
    order_status = _get_order_status()
    order_status.create(op.get_bind())
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('order',
                    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
                    sa.Column('status', order_status, nullable=True),
                    sa.PrimaryKeyConstraint('id', name=op.f('order_pkey'))
                    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_table('order')
    # ### end Alembic commands ###
    _get_order_status().drop(op.get_bind())

Alteration

That's great. We finally managed to make a correct migration. But the application's development is ongoing, and we needed to add a new order status. And we can easily do it:

class OrderStatusEnum(Enum):
    WAITING_FOR_WORKER = 'WAITING_FOR_WORKER'
    IN_PROGRESS = 'IN_PROGRESS'
    DONE = 'DONE'
    CANCELED = 'CANCELED'

Solution I (plain sql)

Let's try to run the autogenerate migration now. And what happened? That's right, nothing! alembic simply ignores the enum changes and the migration is not automatically generated. So what to do about it? For example, you can add plain sql to the migration:

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # ### end Alembic commands ###
    op.execute("ALTER TYPE order_status_enum ADD VALUE 'CANCELED'")


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # ### end Alembic commands ###
    query_str = f"""
            ALTER TYPE order_status_enum RENAME TO order_status_enum_old;
            CREATE TYPE order_status_enum AS ENUM('WAITING_FOR_WORKER', 'IN_PROGRESS', 'DONE');
            ALTER TABLE "order" ALTER status DROP DEFAULT;
            ALTER TABLE "order" ALTER COLUMN status TYPE order_status_enum USING status::text::order_status_enum;
            DROP TYPE order_status_enum_old;
            """
    for q in query_str.strip().split(';')[:-1]:
        op.execute(q)

Solution II (function)

This is a perfectly acceptable solution, but not too convenient, because you will have to constantly remember about this problem and copy the code from old migrations. What should we do? The first idea that came to mind is to simply wrap this in a function:

from typing import Iterable, Sequence
from alembic import op


def set_enum_values(enum_name: str, new_values: Iterable[str], references: Iterable[Sequence[str]]):
    """

    @param enum_name: Enum system name
    @param new_values: New enum values
    @param references: References to enum in models

    Example:
        set_enum_values('promo_type_enum', (
            'BEST_OFFER',
            'NEW_PRODUCT',
            'NO_PROMOTION',
        ), [('advertisement_sale_package', 'promo_type')])
    """
    query_str = f"""
            ALTER TYPE {enum_name} RENAME TO {enum_name}_old;
            CREATE TYPE {enum_name} AS ENUM({', '.join(f"'{value}'" for value in new_values)});
            """
    for table_name, column_name in references:
        query_str += f"""
            ALTER TABLE "{table_name}" ALTER {column_name} DROP DEFAULT;
            ALTER TABLE "{table_name}" ALTER COLUMN {column_name} TYPE {enum_name} USING {column_name}::text::{enum_name};
        """
    query_str += f"""DROP TYPE {enum_name}_old;"""
    for q in query_str.split(';')[:-1]:
        op.execute(q)

Then downgrade can be rewritten as follows:

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # ### end Alembic commands ###
    set_enum_values('order_status_enum', (
        'WAITING_FOR_WORKER',
        'IN_PROGRESS',
        'DONE'
    ), [('order', 'status')])

Solution III (library)

But it's still worth keeping in mind and adding it by hand.
For this reason @RustyGuard, with my support, create the wonderful library
alembic-postgresql-enum.
It completely solves the problems with enum migrations.

Install:

pip install alembic-postgresql-enum

To use it, you only need to add one line to env.py:

# env.py
import alembic_postgresql_enum

Now you can finally forget about most of the problems associated with enum migrations.
I urge everyone who uses postgresql together with sqlalchemy to try it and write about problems encountered.

DEV Community: Alexandrov Roman