How to use Enum with SQLAlchemy and Alembic?

Question

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.

Answer 1

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)

Answer 2

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

Answer 3

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.

Answer 4

This and related StackOverflow threads resort to PostgreSQL dialect-specific typing. However, generic support may be easily achieved in an Alembic migration as follows.

First, import the Python enum, the SQLAlchemy Enum, and your SQLAlchemy declarative base wherever you're going to declare your custom SQLAlchemy Enum column type.

import enum
from sqlalchemy import Enum
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()

Let's take OP's original Python enumerated class:

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

Now we create a SQLAlchemy Enum instantiation:

PostStatusType: Enum = Enum(
    PostStatus,
    name="post_status_type",
    create_constraint=True,
    metadata=Base.metadata,
    validate_strings=True,
)

When you run your Alembic alembic revision --autogenerate -m "Revision Notes" and try to apply the revision with alembic upgrade head, you'll likely get an error about the type not existing. For example:

...
sqlalchemy.exc.ProgrammingError: (psycopg2.errors.UndefinedObject) type "post_status_type" does not exist
LINE 10:  post_status post_status_type NOT NULL,
...

To fix this, import your SQLAlchemy Enum class and add the following to your upgrade() and downgrade() functions in the Alembic autogenerated revision script.

from myproject.database import PostStatusType
...
def upgrade() -> None:
    PostStatusType.create(op.get_bind(), checkfirst=True)
    ... the remainder of the autogen code...
def downgrade() -> None:
    ...the autogen code...
    PostStatusType.drop(op.get_bind(), checkfirst=True)

Finally, be sure to update the auto-generated sa.Column() declaration in the table(s) using the enumerated type to simply reference the SQLAlchemy Enum type instead of using Alembic's attempt to re-declare it. For example in def upgrade() -> None:

op.create_table(
    "my_table",
    sa.Column(
        "post_status",
        PostStatusType,
        nullable=False,
    ),
)