alembic-postgresql-enum

Alembic autogenerate support for creation, alteration and deletion of enums

Alembic will now automatically:

• Create enums that currently are not in postgres schema
• Remove/add/alter enum values
• Reorder enum values
• Delete unused enums from schema

If you are curious to know about analogs and reasons for this library to exist see alternatives and motivation

Usage

Install library:

pip install alembic-postgresql-enum

Add the line:

# env.py
import alembic_postgresql_enum
...

To the top of your migrations/env.py file.

Configuration

You can configure this extension to disable parts of it, or to enable some feature flags

To do so you need to call set_configuration function after the import:

import alembic_postgresql_enum

alembic_postgresql_enum.set_configuration(
    alembic_postgresql_enum.Config(
        add_type_ignore=True,
    )
)

Features

• Creation of enums
• Deletion of unreferenced enums
• Creation of new enum values
• Deletion of enums values
• Renaming of enum values
• Omitting managing enums

Creation of enum

When table is created

class MyEnum(enum.Enum):
    one = 1
    two = 2
    three = 3


class ExampleTable(BaseModel):
    test_field = Column(Integer, primary_key=True, autoincrement=False)
    enum_field = Column(postgresql.ENUM(MyEnum)) 

This code will generate migration given below:

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # this line is generated by our library
    sa.Enum('one', 'two', 'three', name='myenum').create(op.get_bind())
    op.create_table('example_table',
    sa.Column('test_field', sa.Integer(), nullable=False),
    # create_type=False argument is now present on postgresql.ENUM as library takes care of enum creation
    sa.Column('enum_field', postgresql.ENUM('one', 'two', 'three', name='myenum', create_type=False), nullable=True),
    sa.PrimaryKeyConstraint('test_field')
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # drop_table does not drop enum by alembic
    op.drop_table('example_table')
    # It is dropped by us
    sa.Enum('one', 'two', 'three', name='myenum').drop(op.get_bind())
    # ### end Alembic commands ###

When column is added

class MyEnum(enum.Enum):
    one = 1
    two = 2
    three = 3


class ExampleTable(BaseModel):
    test_field = Column(Integer, primary_key=True, autoincrement=False)
    # this column has just been added
    enum_field = Column(postgresql.ENUM(MyEnum)) 

This code will generate migration given below:

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    # this line is generated by our library
    sa.Enum('one', 'two', 'three', name='myenum').create(op.get_bind())
    # create_type=False argument is now present on postgresql.ENUM as library takes care of enum creation
    op.add_column('example_table', sa.Column('enum_field', postgresql.ENUM('one', 'two', 'three', name='myenum', create_type=False), nullable=False))
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('example_table', 'enum_field')
    # enum is explicitly dropped as it is no longer used
    sa.Enum('one', 'two', 'three', name='myenum').drop(op.get_bind())
    # ### end Alembic commands ###

Deletion of unreferenced enum

If enum is defined in postgres schema, but its mentions removed from code - It will be automatically removed

class ExampleTable(BaseModel):
    test_field = Column(Integer, primary_key=True, autoincrement=False)
    # enum_field is removed from table

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('example_table', 'enum_field')
    sa.Enum('one', 'two', 'four', name='myenum').drop(op.get_bind())
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    sa.Enum('one', 'two', 'four', name='myenum').create(op.get_bind())
    op.add_column('example_table', sa.Column('enum_field', postgresql.ENUM('one', 'two', 'four', name='myenum', create_type=False), autoincrement=False, nullable=True))
    # ### end Alembic commands ###

Creation of new enum values

If new enum value is defined sync_enum_values function call will be added to migration to account for it

class MyEnum(enum.Enum):
    one = 1
    two = 2
    three = 3
    four = 4 # New enum value

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.sync_enum_values(
        enum_schema='public', 
        enum_name='myenum', 
        new_values=['one', 'two', 'three', 'four'], 
        affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')],
        enum_values_to_rename=[],
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.sync_enum_values(
        enum_schema='public', 
        enum_name='myenum', 
        new_values=['one', 'two', 'three'], 
        affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')],
        enum_values_to_rename=[],
    )
    # ### end Alembic commands ###

Deletion of enums values

If enum value is removed it also will be detected

class MyEnum(enum.Enum):
    one = 1
    two = 2
    # three = 3 removed

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.sync_enum_values(
        enum_schema='public', 
        enum_name='myenum', 
        new_values=['one', 'two'], 
        affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')],
        enum_values_to_rename=[],
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.sync_enum_values(
        enum_schema='public', 
        enum_name='myenum', 
        new_values=['one', 'two', 'three'], 
        affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')],
        enum_values_to_rename=[],
    )
    # ### end Alembic commands ###

Rename enum value

In this case you must manually edit migration

class MyEnum(enum.Enum):
    one = 1
    two = 2
    three = 3 # renamed from `tree`

This code will generate this migration:

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.sync_enum_values(
        enum_schema='public', 
        enum_name='myenum', 
        new_values=['one', 'two', 'three'], 
        affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')],
        enum_values_to_rename=[],
    )
    # ### end Alembic commands ###


def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.sync_enum_values(
        enum_schema='public', 
        enum_name='myenum', 
        new_values=['one', 'two', 'tree'], 
        affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')],
        enum_values_to_rename=[],
    )
    # ### end Alembic commands ###

This migration will cause problems with existing rows that references MyEnum

So adjust migration like that

def upgrade():
    op.sync_enum_values(
        enum_schema='public', 
        enum_name='myenum', 
        new_values=['one', 'two', 'three'], 
        affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')],
        enum_values_to_rename=[('tree', 'three')],
    )


def downgrade():
    op.sync_enum_values(
        enum_schema='public', 
        enum_name='myenum', 
        new_values=['one', 'two', 'tree'], 
        affected_columns=[TableReference(table_schema='public', table_name='example_table', column_name='enum_field')],
        enum_values_to_rename=[('three', 'tree')],
    )

Do not forget to switch places old and new values for downgrade

All defaults in postgres will be renamed automatically as well

Omitting managing enums

If configured include_name function returns False given enum will be not managed.

import alembic_postgresql_enum

def include_name(name: str) -> bool:
    return name not in ['enum-to-ignore', 'some-internal-enum']

alembic_postgresql_enum.set_configuration(
    alembic_postgresql_enum.Config(
        include_name=include_name,
    )
)

Feature is similar to sqlalchemy feature for tables