Give AlbumentationsX a star on GitHub — it powers this leaderboard

Star on GitHub

alembic-postgresql-enum

Alembic autogenerate support for creation, alteration and deletion of enums

Rank: #3146Downloads: 1,553,345 (30 days)

Description

alembic-postgresql-enum

<img src="https://img.shields.io/pypi/pyversions/alembic-postgresql-enum"> <img src="https://img.shields.io/pypi/v/alembic-postgresql-enum"> <img src="https://img.shields.io/pypi/l/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.

This import will affect newly generated migrations. To try it out you can edit some enums in your schema and then run alembic revision --autogenerate

Features

Creation of enum<a id="creation-of-enum"></a>

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<a id="deletion-of-unreferenced-enum"></a>

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 ###

Detection of enum values changes<a id="detection-of-enum-values-changes"></a>

Can be disabled with detect_enum_values_changes configuration flag turned off

Creation of new enum values<a id="creation-of-new-enum-values"></a>

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<a id="deletion-of-enums-values"></a>

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<a id="rename-enum-value"></a>

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

Partial index preservation<a id="partial-index-preservation"></a>

When modifying enum values, partial indexes that reference the enum type are preserved via dropping and recreating. This is particularly important for indexes with WHERE clauses that use enum comparisons. Depending on the size and complexity of the index this might impact the speed and locking nature of the schema migration.

Note: For alembic's offline mode support, partial index detection happens during migration generation time. The detected indexes are then passed to the migration as a parameter. This ensures that offline migrations can execute without needing database access