Observation
In order to alleviate the pain when migrating, I always use non-native enums even with PostgreSQL.
Non-native enums are just strings with constraints, if you edit an enum, there are only three scenarios:
- Rename enum values
- Remove enum values
- Add enum values.
For migrations, 2 & 3 are a pair. This is understandable: if you upgrade for adding, then you have to remove when downgrading and vice versa. So let's categorize them as two types.
Implementation
If it's renaming, normally I'll break it into three steps:
- Drop old constraints
- Update the rows' old value to the new value
- Create new constraints
In alembic, this is done by:
def update_enum(
table, column, enum_class_name, target_values, olds_to_remove, news_to_add
):
op.drop_constraint(f"ck_{table}_{enum_class_name}", table)
for sql in update_enum_sqls(table, column, olds_to_remove, news_to_add):
op.execute(sql)
op.create_check_constraint(
enum_class_name, table, sa.sql.column(column).in_(target_values)
)
Let's forget update_enum_sqls first and just use it as a SQL generator.
If it's removing, well there are still three steps:
- Drop old constraints
- Delete rows with the old value
- Create new constraints
So basically only the update_enum_sqls might behave differently.
If it's adding, only two steps:
- Drop old constraints
- Create new constraints
Still, we can ignore the update_enum_sqls.
So how to implement it? Not that hard...
def update_enum_sql(table, column, old_value, new_value):
if new_value is not None:
return f"UPDATE {table} SET {column} = '{new_value}' where {column} = '{old_value}'"
else:
return f"DELETE FROM {table} where {column} = '{old_value}'"
def update_enum_sqls(table, column, olds_to_remove, news_to_add):
if len(olds_to_remove) != len(news_to_add):
raise NotImplementedError
return [
update_enum_sql(table, column, old, new)
for old, new in zip(olds_to_remove, news_to_add)
]
Example
Since we prepared the ingredients, let's apply:
def upgrade():
# rename enum
update_enum(
"my_table",
"my_enum",
"myenumclassname",
["NEW", "ENUM", "VALUES"],
["OLD"],
["NEW"],
)
# add enum
update_enum(
"my_table",
"my_enum",
"myenumclassname",
["NEW", "ENUM", "VALUES"],
[],
[],
)
def downgrade():
# remove enum
update_enum(
"my_table",
"my_enum",
"myenumclassname",
["ENUM", "VALUES"],
["NEW"],
[None], # this will delete rows with "NEW", USE WITH CARE!!!
)
# edit enum
update_enum(
"my_table",
"my_enum",
"myenumclassname",
["OLD", "ENUM", "VALUES"],
["NEW"],
["OLD"],
)
The code above can also be found on gist.
