How to use Enum with SQLAlchemy and Alembic?

user_1275294 0
Viewed 26151

Here's my Post model:

class Post(Base):
    __tablename__ = 'posts'

    title = db.Column(db.String(120), nullable=False)
    description = db.Column(db.String(2048), nullable=False)

I'd like to add Enum status to it. So, I've created a new Enum:

import enum

class PostStatus(enum.Enum):
    DRAFT='draft'
    APPROVE='approve'
    PUBLISHED='published'

And added a new field to model:

class Post(Base):
    ...
    status = db.Column(db.Enum(PostStatus), nullable=False, default=PostStatus.DRAFT.value, server_default=PostStatus.DRAFT.value)

After doing FLASK_APP=server.py flask db migrate, a such migration was generated:

def upgrade():
    op.add_column('posts', sa.Column('status', sa.Enum('DRAFT', 'APPROVE', 'PUBLISHED', name='poststatus'), server_default='draft', nullable=False))

After trying to upgrade DB, I'm getting:

sqlalchemy.exc.ProgrammingError: (psycopg2.ProgrammingError) type "poststatus" does not exist
LINE 1: ALTER TABLE posts ADD COLUMN status poststatus DEFAULT 'draf...
                                            ^
 [SQL: "ALTER TABLE posts ADD COLUMN status poststatus DEFAULT 'draft' NOT NULL"]
  1. Why type poststatus was not created on DB-level automatically? In the similar migration it was.
  2. How to specify server_default option properly? I need both ORM-level defaults and DB-level ones, because I'm altering existing rows, so ORM defaults are not applied.
  3. Why real values in DB are 'DRAFT', 'APPROVE', 'PUBLISHED', but not draft, etc? I supposed there should be ENUM values, not names.

Thank you in advance.

alembic flask-sqlalchemy sqlalchemy python
4 Answers

Why real values in DB are 'DRAFT', 'APPROVE', 'PUBLISHED', but not draft, etc? I supposed there should be ENUM values, not names.

As Peter Bašista's already mentioned SQLAlchemy uses the enum names (DRAFT, APPROVE, PUBLISHED) in the database. I assume that was done because the enum values ("draft", "approve", ...) can be arbitrary types in Python and they are not guaranteed to be unique (unless @unique is used).

However since SQLAlchemy 1.2.3 the Enum class accepts a parameter values_callable which can be used to store enum values in the database:

    status = db.Column(
        db.Enum(PostStatus, values_callable=lambda obj: [e.value for e in obj]),
        nullable=False,
        default=PostStatus.DRAFT.value,
        server_default=PostStatus.DRAFT.value
    )

Why type poststatus was not created on DB-level automatically? In the similar migration it was.

I think basically you are hitting a limitation of alembic: It won't handle enums on PostgreSQL correctly in some cases. I suspect the main issue in your case is Autogenerate doesn't correctly handle postgresql enums #278.

I noticed that the type is created correctly if I use alembic.op.create_table so my workaround is basically:

enum_type = SQLEnum(PostStatus, values_callable=lambda enum: [e.value for e in enum])
op.create_table(
    '_dummy',
    sa.Column('id', Integer, primary_key=True),
    sa.Column('status', enum_type)
)
op.drop_table('_dummy')
c_status = Column('status', enum_type, nullable=False)
add_column('posts', c_status)
user_1385260

Use the following function example in case you are using PostgreSQL:

from sqlalchemy.dialects import postgresql
from ... import PostStatus
from alembic import op
import sqlalchemy as sa


def upgrade():
    post_status = postgresql.ENUM(PostStatus, name="status")
    post_status.create(op.get_bind(), checkfirst=True)
    op.add_column('posts', sa.Column('status',  post_status))


def downgrade():
    post_status = postgresql.ENUM(PostStatus, name="status")
    post_status.drop(op.get_bind())
user_8956590

I can only answer the third part of your question.

The documentation for the Enum type in SQLAlchemy states that:

Above, the string names of each element, e.g. “one”, “two”, “three”, are persisted to the database; the values of the Python Enum, here indicated as integers, are not used; the value of each enum can therefore be any kind of Python object whether or not it is persistable.

So, it is by SQLAlchemy design that Enum names, not values are persisted into the database.

user_10184470

From official docs: https://docs.python.org/3/library/enum.html#creating-an-enum

import enum

class PostStatus(enum.Enum):
    DRAFT = 0
    APPROVE = 1
    PUBLISHED = 2

According to this:

class Post(Base):
    ...
    status = db.Column(db.Integer(), nullable=False, default=PostStatus.DRAFT.value, server_default=PostStatus.DRAFT.value)

1) PostStatus is not a DB-model, it's just a class which contains status ids;

2) it's OK

3) you don't have to store status strings in DB, you better use ids instead

user_36529490
Related