Использование PostgreSQL Enum в SQLAlchemy. И проблемы с ним

Начнем с начала. В 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. Пишите если есть что добавить и пробуйте библиотеку.