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
I had the same issue trying to migrate a column type to another. I use the following requirements:
Alembic==0.9.4
SQLAlchemy==1.1.12
You can provide the argument postgresql_using
as a kwarg of alembic.op.alter_column
.
from alembic import op
import sqlalchemy as types
op.alter_column(
table_name='my_table',
column_name='my_column',
type_=types.NewType,
# allows to use postgresql USING
postgresql_using="my_column::PostgesEquivalentOfNewType",
)
I hope it can help.
I decided to try to follow the postgres approach as directly as possible and came up with the following migration.
from alembic import op
import sqlalchemy as sa
old_options = ('nonexistent_executable', 'signal', 'success', 'timed_out')
new_options = sorted(old_options + ('output_limit_exceeded',))
old_type = sa.Enum(*old_options, name='status')
new_type = sa.Enum(*new_options, name='status')
tmp_type = sa.Enum(*new_options, name='_status')
tcr = sa.sql.table('testcaseresult',
sa.Column('status', new_type, nullable=False))
def upgrade():
# Create a tempoary "_status" type, convert and drop the "old" type
tmp_type.create(op.get_bind(), checkfirst=False)
op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE _status'
' USING status::text::_status')
old_type.drop(op.get_bind(), checkfirst=False)
# Create and convert to the "new" status type
new_type.create(op.get_bind(), checkfirst=False)
op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE status'
' USING status::text::status')
tmp_type.drop(op.get_bind(), checkfirst=False)
def downgrade():
# Convert 'output_limit_exceeded' status into 'timed_out'
op.execute(tcr.update().where(tcr.c.status==u'output_limit_exceeded')
.values(status='timed_out'))
# Create a tempoary "_status" type, convert and drop the "new" type
tmp_type.create(op.get_bind(), checkfirst=False)
op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE _status'
' USING status::text::_status')
new_type.drop(op.get_bind(), checkfirst=False)
# Create and convert to the "old" status type
old_type.create(op.get_bind(), checkfirst=False)
op.execute('ALTER TABLE testcaseresult ALTER COLUMN status TYPE status'
' USING status::text::status')
tmp_type.drop(op.get_bind(), checkfirst=False)
It appears that alembic has no direct support for the USING
statement in its alter_table
method.
I used a bit simpler approach with less steps than the accepted answer, which I based this on. In this example I will pretend the enum in question is called 'status_enum', because in the accepted answer the use of 'status' for both the column and enum confused me.
from alembic import op
import sqlalchemy as sa
name = 'status_enum'
tmp_name = 'tmp_' + name
old_options = ('nonexistent_executable', 'signal', 'success', 'timed_out')
new_options = sorted(old_options + ('output_limit_exceeded',))
new_type = sa.Enum(*new_options, name=name)
old_type = sa.Enum(*old_options, name=name)
tcr = sa.sql.table('testcaseresult',
sa.Column('status', new_type, nullable=False))
def upgrade():
op.execute('ALTER TYPE ' + name + ' RENAME TO ' + tmp_name)
new_type.create(op.get_bind())
op.execute('ALTER TABLE testcaseresult ALTER COLUMN status ' +
'TYPE ' + name + ' USING status::text::' + name)
op.execute('DROP TYPE ' + tmp_name)
def downgrade():
# Convert 'output_limit_exceeded' status into 'timed_out'
op.execute(tcr.update().where(tcr.c.status=='output_limit_exceeded')
.values(status='timed_out'))
op.execute('ALTER TYPE ' + name + ' RENAME TO ' + tmp_name)
old_type.create(op.get_bind())
op.execute('ALTER TABLE testcaseresult ALTER COLUMN status ' +
'TYPE ' + name + ' USING status::text::' + name)
op.execute('DROP TYPE ' + tmp_name)
Since I got conversion errors and problems with default values, I wrote an even more generalised answer based on the accepted one:
def replace_enum_values(
name: str,
old: [str],
new: [str],
modify: [(str, str, str)]
):
"""
Replaces an enum's list of values.
Args:
name: Name of the enum
new: New list of values
old: Old list of values
modify: List of tuples of table name
and column to modify (which actively use the enum).
Assumes each column has a default val.
"""
connection = op.get_bind()
tmp_name = "{}_tmp".format(name)
# Rename old type
op.execute(
"ALTER TYPE {} RENAME TO {};"
.format(name, tmp_name)
)
# Create new type
lsl = sa.Enum(*new, name=name)
lsl.create(connection)
# Replace all usages
for (table, column) in modify:
# Get default to re-set later
default_typed = connection.execute(
"SELECT column_default "
"FROM information_schema.columns "
"WHERE table_name='{table}' "
"AND column_name='{column}';"
.format(table=table, column=column)
).first()[0] # type: str
# Is bracketed already
default = default_typed[:default_typed.index("::")]
# Set all now invalid values to default
connection.execute(
"UPDATE {table} "
"SET {column}={default} "
"WHERE {column} NOT IN {allowed};"
.format(
table=table,
column=column,
# Invalid: What isn't contained in both new and old
# Can't just remove what's not in new because we get
# a type error
allowed=tuple(set(old).intersection(set(new))),
default=default
)
)
op.execute(
"ALTER TABLE {table} "
# Default needs to be dropped first
"ALTER COLUMN {column} DROP DEFAULT,"
# Replace the tpye
"ALTER COLUMN {column} TYPE {enum_name} USING {column}::text::{enum_name},"
# Reset default
"ALTER COLUMN {column} SET DEFAULT {default};"
.format(
table=table,
column=column,
enum_name=name,
default=default
)
)
# Remove old type
op.execute("DROP TYPE {};".format(tmp_name))
This can be called from upgrade / downgrade as such:
replace_enum_values(
name='enum_name',
new=["A", "B"],
old=["A", "C"],
modify=[('some_table', 'some_column')]
)
All invalidated values will be set to server_default.