Postgresql – How to Change a Column from ENUM to INTEGER (Foreign Key) Using Alembic?

This is kind of messy but I had to make the column nullable with one alter before I could make the subsequent alter to change the type and set all the values to NULL. You can actually migrate the enum to hardcoded ints if you know the mapping and use a CASE statement which you said you didn’t need but I thought was pretty cool.

• create dummy data
• change column to null
• change column’s type to int
• create foreign key
• set a single fk value
• check that relationship is working

import os

from sqlalchemy import (
    create_engine,
    Column,
    Integer,
    String,
    ForeignKey,
    DateTime,
)
import alembic
from sqlalchemy.orm import Mapped, aliased
from sqlalchemy.orm import mapped_column
from sqlalchemy.sql import (
    text,
    select,
    func,
    and_,
    or_,
    update,
    insert,
    delete,
)
from sqlalchemy.orm import (
    DeclarativeBase,
    Session,
    relationship,
    selectinload,
    joinedload,
)
from typing import List
from sqlalchemy import event
from sqlalchemy.schema import CreateSchema
from sqlalchemy.dialects import postgresql
from alembic.migration import MigrationContext
from alembic.operations import Operations


def get_engine(env):
    return create_engine(f"postgresql://{env['DB_USER']}:{env['DB_PASSWORD']}@{env['DB_HOST']}:{env['DB_PORT']}/{env['DB_NAME']}", echo=True)



def get_meta(with_enum=True):

    class Base(DeclarativeBase):
        pass

    # I am sure there is a smarter way to do this but for now
    # I just redefine the entire metadata after the migration.
    if with_enum:
        class Serie(Base):
            __tablename__ = 'series'
            id: Mapped[int] = mapped_column(primary_key=True)
            name: Mapped[str] = mapped_column()
            serie_aluno: Mapped[postgresql.ENUM] = mapped_column(postgresql.ENUM('6', '7', '8', '9', name='serie_enum'), nullable=False)
    else:
        class Serie(Base):
            __tablename__ = 'series'
            id: Mapped[int] = mapped_column(primary_key=True)
            name: Mapped[str] = mapped_column()
            serie_aluno: Mapped[int] = mapped_column(ForeignKey('alunos.id'))
            aluno: Mapped["Aluno"] = relationship("Aluno")

    class Aluno(Base):
        __tablename__ = 'alunos'
        id: Mapped[int] = mapped_column(primary_key=True)
        name: Mapped[str]

    return Base, Serie, Aluno


def main(env):
    engine = get_engine(env)

    Base, Serie, Aluno = get_meta(with_enum=True)
    with engine.connect() as conn, conn.begin():

        Base.metadata.create_all(conn)

    with Session(engine) as db:
        db.add(Serie(name='serie1', serie_aluno='6'))
        db.add(Serie(name='serie2', serie_aluno='9'))
        db.add(Aluno(name='aluno1'))
        db.commit()

    with engine.connect() as conn:
        with conn.begin() as trans:
            context = MigrationContext.configure(conn)
            op = Operations(context)
            #whens = ' '.join(f"WHEN serie_aluno='{v}' THEN {v}" for v in ('6', '7', '8', '9'))
            #postgresql_using=f"CASE {whens} END"
            # Change from nullable = False to nullable = True
            op.alter_column(Serie.__tablename__, "serie_aluno", nullable=True, existing_nullable=False)
            # Alter column and set values to NULL
            op.alter_column(Serie.__tablename__, "serie_aluno", existing_type=postgresql.ENUM, type_=Integer, postgresql_using="NULL")
            op.create_foreign_key(None, 'series', 'alunos', ['serie_aluno'], ['id'])
            trans.commit()

    Base, Serie, Aluno = get_meta(with_enum=False)

    with Session(engine) as db:
        db.execute(text("UPDATE series SET serie_aluno = 1 WHERE name = 'serie1'"))
        db.flush()
        assert db.scalar(select(func.count(Serie.id)).where(Serie.serie_aluno != None)) == 1



if __name__ == '__main__':
    main(os.environ)