Merge "Add setting server_default in alter_enum(_add_value)"

This commit is contained in:
Zuul 2018-03-17 01:53:57 +00:00 committed by Gerrit Code Review
commit fcdf3c1d37
1 changed files with 19 additions and 7 deletions

View File

@ -112,7 +112,7 @@ 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, enum, nullable):
def alter_enum_add_value(table, column, enum, nullable, server_default=None):
'''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.
@ -121,6 +121,7 @@ def alter_enum_add_value(table, column, enum, nullable):
:param column: column name
:param enum: sqlalchemy Enum with updated values
:param nullable: existing nullable for column.
:param server_default: existing or new server_default for the column
'''
bind = op.get_bind()
@ -129,25 +130,34 @@ def alter_enum_add_value(table, column, enum, nullable):
values = {'name': enum.name,
'values': ", ".join("'" + i + "'" for i in enum.enums),
'column': column,
'table': table}
'table': table,
'server_default': server_default}
if server_default is not None:
op.execute("ALTER TABLE %(table)s ALTER COLUMN %(column)s"
" DROP DEFAULT" % values)
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)
if server_default is not None:
op.execute("ALTER TABLE %(table)s ALTER COLUMN %(column)s"
" SET DEFAULT '%(server_default)s'" % values)
op.execute("DROP TYPE old_%(name)s" % values)
else:
op.alter_column(table, column, type_=enum,
existing_nullable=nullable)
existing_nullable=nullable,
server_default=server_default)
def alter_enum(table, column, enum_type, nullable, do_drop=True,
def alter_enum(table, column, enum_type, nullable,
server_default=None, do_drop=True,
do_rename=True, do_create=True):
"""Alter a enum type column.
Set the do_xx parameters only when the modified enum type
is used by multiple columns. Else don't provide these
parameters.
:param server_default: existing or new server_default for the column
:param do_drop: set to False when modified column is
not the last one use this enum
:param do_rename: set to False when modified column is
@ -167,7 +177,8 @@ def alter_enum(table, column, enum_type, nullable, do_drop=True,
enum_type.create(bind, checkfirst=False)
op.execute("ALTER TABLE %(table)s RENAME COLUMN %(column)s TO "
"old_%(column)s" % values)
op.add_column(table, sa.Column(column, enum_type, nullable=nullable))
op.add_column(table, sa.Column(column, enum_type, nullable=nullable,
server_default=server_default))
op.execute("UPDATE %(table)s SET %(column)s = " # nosec
"old_%(column)s::text::%(name)s" % values)
op.execute("ALTER TABLE %(table)s DROP COLUMN old_%(column)s" % values)
@ -175,7 +186,8 @@ def alter_enum(table, column, enum_type, nullable, do_drop=True,
op.execute("DROP TYPE old_%(name)s" % values)
else:
op.alter_column(table, column, type_=enum_type,
existing_nullable=nullable)
existing_nullable=nullable,
server_default=server_default)
def create_table_if_not_exist_psql(table_name, values):