Fix alter_enum_add_value

Using alter_enum_add_value for PostgreSQL hit error with running
inside migration.

Update code with operations that can be executed. As this implementation
do not drop columns like alter_enum, it can be used in expand branch
as well.

Change-Id: I56d0731fff317b7be73e8237fa897c1d44331bcd
This commit is contained in:
Ann Kamyshnikova 2017-07-05 16:19:51 +04:00 committed by Ann Taraday
parent df762b0d46
commit 42fbe0e0b5

View File

@ -107,14 +107,13 @@ def rename_table_if_exists(old_table_name, new_table_name):
op.rename_table(old_table_name, new_table_name)
def alter_enum_add_value(table, column, new_value, enum, nullable):
def alter_enum_add_value(table, column, enum, nullable):
'''If we need to expand Enum values for some column - for PostgreSQL this
can be done with ALTER TYPE function. For MySQL, it can be done with
ordinary alembic alter_column function.
:param table:table name
:param column: column name
:param new_value: value that we want to add for Enum
:param enum: sqlalchemy Enum with updated values
:param nullable: existing nullable for column.
'''
@ -122,9 +121,15 @@ def alter_enum_add_value(table, column, new_value, enum, nullable):
bind = op.get_bind()
engine = bind.engine
if engine.name == 'postgresql':
values = {'value': new_value,
'name': enum.name}
op.execute("ALTER TYPE %(name)s ADD VALUE '%(value)s'" % values)
values = {'name': enum.name,
'values': ", ".join("'" + i + "'" for i in enum.enums),
'column': column,
'table': table}
op.execute("ALTER TYPE %(name)s rename to old_%(name)s" % values)
op.execute("CREATE TYPE %(name)s AS enum (%(values)s)" % values)
op.execute("ALTER TABLE %(table)s ALTER COLUMN %(column)s TYPE "
"%(name)s USING %(column)s::text::%(name)s " % values)
op.execute("DROP TYPE old_%(name)s" % values)
else:
op.alter_column(table, column, type_=enum,
existing_nullable=nullable)