Using PostgreSQL Enum in SQLAlchemy. And problems with it
Start over. In postgres, an enum is a static ordered set of values.
In the course of using postgresql with sqlalchemy, I encountered a number of features, which are poorly documentedwhich I would like to talk about.
This article will consist of 3 parts:
Basic use cases
Imagine we have an order model
from sqlalchemy import Integer
from app.database import BaseModel
class Order(BaseModel):
id = Column(Integer, primary_key=True, autoincrement=True)
# ... другие поляMost likely, he will have different statuses. This can be implemented as a simple numeric field, but there are problems with limiting such a field and, in general, interpreting it when fetching from the database. So it’s better to use enum. For example, let’s write a simple enum using the python standard library.
from enum import Enum
class OrderStatusEnum(Enum):
WAITING_FOR_WORKER = 'WAITING_FOR_WORKER'
IN_PROGRESS = 'IN_PROGRESS'
DONE = 'DONE'
And now, to add an enumeration column, you just need to add a similar line
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)
# ... другие поля
About options:
name – the name with which the enum will be created in the database
create_type – automatic creation/deletion (which does not work with automatic generation) of an enumeration in a migration. For this reason, I do not recommend using this option.
Problems with migrations
So, now let’s move on to the section, because of which this article was written in many respects. We wrote the first enum, and if you do not write the create_type=False parameter, then it may seem that everything is fine and works as it should. The migration has been 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"), 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 it is not so. If we do a downgrade, and then try to upgrade again, it turns out that the enum was not deleted in the downgrade. How to fix it? We need to add a line that will remove our enum. It will look like this (for uniformity of the code, I also registered 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())Great. We were finally able to create an enum and make a correct migration. But the development of the application does not stand still, and we needed to add a new order status. And we can easily do this:
class OrderStatusEnum(Enum):
WAITING_FOR_WORKER = 'WAITING_FOR_WORKER'
IN_PROGRESS = 'IN_PROGRESS'
DONE = 'DONE'
CANCELED = 'CANCELED'Let’s try to run the migration auto-generation now. And what happened? That’s right, nothing! alembic simply ignores enum changes and no migration is generated automatically. What to do with it? For example, you can add something like this 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)
This is quite a solution, but it is not very convenient, since you will have to constantly remember about this problem and copy code from old migrations. What to do? The first idea that came up was to just wrap it 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
@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)
Then downgrade can be rewritten like this:
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')])But this is still worth remembering and adding pens. For this reason @RustyGuard created a wonderful library with my support alembic-postgresql-enum. It completely solves problems with enum migrations. I encourage everyone who uses postgresql with sqlalchemy to give it a try and report any problems they encounter.
Advanced usage
Finally, I would like to talk about the case when I had to change the behavior in the default enum. An important detail should be noted here, when generating an enumeration, sqlalchemy takes the name attribute, not the value attribute. And when I needed to write such an enumeration
class DecorationState(Enum):
NONE = 'Без отделки'
MUNICIPAL = 'Муниципальный ремонт'
MODERN = 'Современная отделка'There was a problem, because I needed to filter by text, and it was necessary that the value attributes, not the name, get into the database. Fortunately, sqlalchemy allows you to create custom types, and I successfully solved this problem by writing a class like this:
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 processThis is my experience with enums in sqlalchemy. Write if you have something to add and try library.
