Использование PostgreSQL Enum в SQLAlchemy. И проблемы с ним
Начнем с начала. В postgres enum - статический упорядоченный набор значений.
В ходе использования postgresql совместно с sqlalchemy я столкнулся с рядом особенностей, которые слабо документированы, о которых мне и хотелось бы поведать.
Данная статья будет состоять из 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 позволяет создавать пользовательские типы и я успешно решил эту задачу написав такой класс:
values_enum.py
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. Пишите если есть что добавить и пробуйте библиотеку.