Начнем с начала. В postgres enum - статический упорядоченный набор значений.

В ходе использования postgresql совместно с sqlalchemy я столкнулся с рядом особенностей, которые слабо документированы, о которых мне и хотелось бы поведать.

Данная статья будет состоять из 3 частей:

1. Базовые примеры использования

2. Проблемы с миграциями

3. Продвинутое использование

Базовые примеры использования

Представим, что у нас есть модель заказа:

from sqlalchemy import Integer
from app.database import BaseModel

class Order(BaseModel):
    id = Column(Integer, primary_key=True, autoincrement=True)
    # ... другие поля

Вероятнее всего, у него будут разные статусы. Это можно реализовать как просто числовое поле, но появляются проблемы с ограничением такого поля и в целом интерпритации его при выборке из бд. Поэтому лучше использовать enum. Для примера давайте напишем простейший перечисление используя стандартную библиотеку python.

from enum import Enum

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

И теперь для добавления колонки-перечисления требуется всего лишь добавить подобную строку:

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', create_type=False), nullable=False, default=OrderStatusEnum.WAITING_FOR_WORKER)
    # ... другие поля

Про параметры:

• name - имя, с которым будет создан enum в бд;

• create_type - автоматическое создание/удаление (которое при автоматической генерации не работает) перечисления в миграции. По этой причине не советую пользоваться данным параметром.

Проблемы с миграциями

Итак, теперь переходим к разделу, из-за которого во многом эта статья и была написана. Мы написали первый enum, и если не прописывать параметр create_type=False, то может показаться что все хорошо и работает как надо. Миграция создалась и применилась.

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'), 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 ###

Но это не совсем так. Если мы сделаем downgrade, а потом попробуем вновь сделать upgrade, то выяснится, что в downgrade не удалился enum. Как же это исправить? Необходимо добавить строчку, которая удалит наш enum. Выглядеть это будет следующим образом (для единообразия кода также прописал явное создание).

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())

Отлично. Мы наконец смогли создать enum и сделать корректную миграцию. Но развитие приложения не стоит на месте, и нам понадобилось добавить новый статус заказа. И мы легко можем это сделать:

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

Давайте же попробуем запустить автогенерацию миграции теперь. И что же произошло? Правильно, ничего! alembic просто игнорирует изменения в enum-е и автоматически миграция не генерируется. Что же с этим делать? Например, можно добавить подобное в миграцию

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)

Это вполне себе решение, но оно не слишком удобно, так как об этой проблеме придется постоянно помнить и копировать код из старых миграций. Что же делать? Первое появившаяся идея - просто обернуть это в функцию:

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
    @param new_values: Новые значения enum
    @param references: Упоминания enum в моделях

    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)

Тогда downgrade можно переписать следующим образом:

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')])

Но об этом по прежнему стоит помнить и добавлять ручками. По этой причине @RustyGuard при моей поддержке создал замечательную библиотеку alembic-postgresql-enum. Она полностью решает проблемы с миграциями enum-ов. Призываю всех кто использует postgresql совместно с sqlalchemy попробовать ее и писать о встреченных проблемах.

Продвинутое использование

Напоследок, хотел бы рассказать о случае, когда мне пришлось менять поведение в enum по умолчанию. Тут следует отметить важную деталь, при генерации перечисления sqlalchemy берет атрибут name, а не value. И когда мне понадобилось написать подобное перечисление:

class DecorationState(Enum):
    NONE = 'Без отделки'
    MUNICIPAL = 'Муниципальный ремонт'
    MODERN = 'Современная отделка'

Возникла проблема, так как мне требовалась фильтрация по тексту, и необходимо было чтобы в базу данных попадали атрибуты value, а не name. К счастью sqlalchemy позволяет создавать пользовательские типы и я успешно решил эту задачу написав такой класс:

from copy import copy
from enum import Enum

import sqlalchemy.types as types


class ValuesEnum(types.TypeDecorator):
    impl = types.Enum

    cache_ok = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        objects = copy(self._object_lookup)
        objects.pop(None)
        self._reversed_object_lookup = {v.value: v for v in objects.values()}
        self._reversed_object_lookup[None] = None

    def process_bind_param(self, value, dialect):
        if isinstance(value, Enum):
            return value.value
        return value

    def _object_value_for_elem(self, value):
        return self._reversed_object_lookup[value]

    def result_processor(self, dialect, coltype):
        def process(value):
            value = self._object_value_for_elem(value)
            return value

        return process

Таков мой опыт работы с enum-ами в sqlalchemy. Пишите если есть что добавить и пробуйте библиотеку.