How can I add an element to an Enum field in an alembic migration when using a version of PostgreSQL older than 9.1 (which adds the ALTER TYPE for enums)? This SO question e
This solution is easy to understand and works really well for both upgrade and downgrade. I've written this answer in more detailed manner.
Let's say our enum_type
looks like this:
enum_type = ('some_value_1', 'some_value_2')
I want to alter enum_type
by adding a new enum, so that it becomes like this:
enum_type = ('some_value_1', 'some_value_2', 'new_value')
This can be done in this way:
from alembic import op
def upgrade():
op.execute("COMMIT")
op.execute("ALTER TYPE enum_type ADD VALUE 'new_value'")
def downgrade():
# Drop 'new_value' from enum_type
op.execute("ALTER TYPE enum_type RENAME TO enum_type_tmp")
op.execute("CREATE TYPE enum_type AS ENUM('some_value_1', 'some_value_1')")
op.execute("DROP TYPE enum_type_tmp")
NOTE: During downgrade, if you're using enum_type
in a table then you can modify the downgrade method as mentioned below:
def downgrade():
# Drop 'new_value' from enum_type
op.execute("UPDATE table_name"
" SET column_name_using_enum_type_value = NULL"
" WHERE column_name_using_enum_type_value = 'new_value'")
op.execute("ALTER TYPE enum_type RENAME TO enum_type_tmp")
op.execute("CREATE TYPE enum_type AS ENUM('some_value_1', 'some_value_1')")
op.execute("ALTER TABLE table_name"
" ALTER COLUMN column_name_using_enum_type_value TYPE enum_type"
" USING column_name_using_enum_type_value::text::enum_type")
op.execute("DROP TYPE enum_type_tmp")
Found another handy method
op.execute('ALTER TYPE enum_type ADD VALUE new_value')
op.execute('ALTER TYPE enum_type ADD VALUE new_value BEFORE old_value')
op.execute('ALTER TYPE enum_type ADD VALUE new_value AFTER old_value')
I needed to move data while migrating types, including deleting some old types, so I figured I'd write up a more general way of doing this based on the (awesome) accepted answer (https://stackoverflow.com/a/14845740/629272). Hopefully this helps someone else in the same boat!
# This migration will move data from one column to two others based on the type
# for a given row, and modify the type of each row.
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
revision = '000000000001'
down_revision = '000000000000'
branch_labels = None
depends_on = None
# This set of options makes up the old type.
example_types_old = (
'EXAMPLE_A',
'EXAMPLE_B',
'EXAMPLE_C',
)
example_type_enum_old = postgresql.ENUM(*example_types_old, name='exampletype')
# This set of options makes up the new type.
example_types_new = (
'EXAMPLE_C',
'EXAMPLE_D',
'EXAMPLE_E',
)
example_type_enum_new = postgresql.ENUM(*example_types_new, name='exampletype')
# This set of options includes everything from the old and new types.
example_types_tmp = set(example_types_old + example_types_new)
example_type_enum_tmp = postgresql.ENUM(*example_types_tmp, name='_exampletype')
# This is a table view from which we can select and update as necessary. This
# only needs to include the relevant columns which are in either the old or new
# version of the table.
examples_view = sa.Table(
# Use the name of the actual table so it is modified in the upgrade and
# downgrade.
'examples',
sa.MetaData(),
sa.Column('id', sa.Integer, primary_key=True),
# Use the _tmp type so all types are usable.
sa.Column('example_type', example_type_enum_tmp),
# This is a column from which the data will be migrated, after which the
# column will be removed.
sa.Column('example_old_column', sa.Integer),
# This is a column to which data from the old column will be added if the
# type is EXAMPLE_A.
sa.Column('example_new_column_a', sa.Integer),
# This is a column to which data from the old column will be added if the
# type is EXAMPLE_B.
sa.Column('example_new_column_b', sa.Integer),
)
def upgrade():
connection = op.get_bind()
# Add the new column to which data will be migrated.
example_new_column_a = sa.Column(
'example_new_column_a',
sa.Integer,
nullable=True
)
op.add_column('examples', example_new_column_a)
# Add the new column to which data will be migrated.
example_new_column_b = sa.Column(
'example_new_column_b',
sa.Integer,
nullable=True
)
op.add_column('examples', example_new_column_b)
# Create the temporary enum and change the example_type column to use the
# temporary enum.
# The USING statement automatically maps the old enum to the temporary one.
example_type_enum_tmp.create(connection, checkfirst=False)
# Change to the temporary type and map from the old type to the temporary
# one.
op.execute('''
ALTER TABLE examples
ALTER COLUMN example_type
TYPE _exampletype
USING example_type::text::_exampletype
''')
# Move data from example_old_column to example_new_column_a and change its
# type to EXAMPLE_D if the type is EXAMPLE_A.
connection.execute(
examples_view.update().where(
examples_view.c.example_type == 'EXAMPLE_A'
).values(
example_type='EXAMPLE_D',
example_new_column_a=examples_view.c.example_old_column,
)
)
# Move data from example_old_column to example_new_column_b and change its
# type to EXAMPLE_E if the type is EXAMPLE_B.
connection.execute(
examples_view.update().where(
examples_view.c.example_type == 'EXAMPLE_B'
).values(
example_type='EXAMPLE_E',
example_new_column_b=examples_view.c.example_old_column,
)
)
# Move any remaining data from example_old_column to example_new_column_a
# and keep its type as EXAMPLE_C.
connection.execute(
examples_view.update().where(
examples_view.c.example_type == 'EXAMPLE_C'
).values(
example_type='EXAMPLE_C',
example_new_column_a=examples_view.c.example_old_column,
)
)
# Delete the old enum now that the data with the old types have been moved.
example_type_enum_old.drop(connection, checkfirst=False)
# Create the new enum and change the example_type column to use the new
# enum.
# The USING statement automatically maps the temporary enum to the new one.
example_type_enum_new.create(connection, checkfirst=False)
op.execute('''
ALTER TABLE examples
ALTER COLUMN example_type
TYPE exampletype
USING example_type::text::exampletype
''')
# Delete the temporary enum.
example_type_enum_tmp.drop(connection, checkfirst=False)
# Remove the old column.
op.drop_column('examples', 'example_old_column')
# The downgrade just performs the opposite of all the upgrade operations but in
# reverse.
def downgrade():
connection = op.get_bind()
example_old_column = sa.Column(
'example_old_column',
sa.Integer,
nullable=True
)
op.add_column('examples', example_old_column)
example_type_enum_tmp.create(connection, checkfirst=False)
op.execute('''
ALTER TABLE examples
ALTER COLUMN example_type
TYPE _exampletype
USING example_type::text::_exampletype
''')
connection.execute(
examples_view.update().where(
examples_view.c.example_type == 'EXAMPLE_C'
).values(
example_type='EXAMPLE_C',
example_old_column=examples_view.c.example_new_column_b,
)
)
connection.execute(
examples_view.update().where(
examples_view.c.example_type == 'EXAMPLE_E'
).values(
example_type='EXAMPLE_B',
example_old_column=examples_view.c.example_new_column_b,
)
)
connection.execute(
examples_view.update().where(
examples_view.c.example_type == 'EXAMPLE_D'
).values(
example_type='EXAMPLE_A',
example_old_column=examples_view.c.example_new_column_a,
)
)
example_type_enum_old.create(connection, checkfirst=False)
op.execute('''
ALTER TABLE examples
ALTER COLUMN example_type
TYPE exampletype
USING example_type::text::exampletype
''')
example_type_enum_tmp.drop(connection, checkfirst=False)
op.drop_column('examples', 'example_new_column_b')
op.drop_column('examples', 'example_new_column_a')
This runs without problems:
from alembic import op
def upgrade():
op.execute("COMMIT")
op.execute("ALTER TYPE enum_type ADD VALUE 'new_value'")
def downgrade():
...
Reference
In straight SQL, this would work for Postgres, if the order of the things in your enum doesn't need to be exactly as above:
ALTER TYPE status ADD value 'output_limit_exceeded' after 'timed_out';
As of Postgres 9.1 adding a new value to an enum can be done with the ALTER TYPE statement. This is complicated by the fact that it cannot be done in a transaction. However this can be worked around by committing alembic's transaction see here.
I actually had problems using the older, more verbose, solution because Postgres could not automatically convert the default for the column.