Como alterar uma coluna de ENUM para INTEGER (chave estrangeira) usando Alembic?

Para quem estiver interessado, a solução foi bem simples, basicamente adicionei o comando sql como meu amigo Erwin sugeriu.

def upgrade() -> None:

  op.execute("ALTER TABLE aluno ALTER COLUMN serie_aluno TYPE integer USING serie_aluno::text::int, ADD CONSTRAINT serie_aluno_serie_fkey FOREIGN KEY (serie_aluno) REFERENCES serie(id_serie);")

def downgrade() -> None:

  op.drop_constraint('serie_aluno_serie_fkey', 'aluno', type_='foreignkey')
  op.alter_column(
        'aluno',
        'serie_aluno',
        type_=sa.Enum('A', 'B', name='turma_enum'),
        postgresql_using="serie_aluno::int::text"
        )

Não há nenhuma conversão definida do seu tipo de enumeração personalizado serie_alunopara integer. (A menos que você crie uma conversão explicitamente.)

Mas tudo pode ser convertido para text. Se todos os valores enum forem representados por literais inteiros válidos, você pode usar textcomo trampolim:

ALTER TABLE aluno ALTER COLUMN serie_aluno TYPE int USING serie_aluno::text::int;

Para então adicionar uma FOREIGN KEYrestrição apontando para uma integercoluna:

ALTER TABLE aluno ADD CONSTRAINT aluno_serie_fkey FOREIGN KEY (serie_aluno) REFERENCES serie(id_serie);

Ou faça isso em um único comando ALTER TABLE :

ALTER TABLE aluno
  ALTER COLUMN serie_aluno TYPE integer USING serie_aluno::text::int
, ADD CONSTRAINT serie_aluno_serie_fkey FOREIGN KEY (serie_aluno) REFERENCES serie(id_serie)
;

violino

Isso é meio bagunçado, mas eu tive que tornar a coluna anulável com uma alteração antes de poder fazer a alteração subsequente para alterar o tipo e definir todos os valores para NULL. Você pode realmente migrar o enum para ints codificados se você souber o mapeamento e usar uma CASEdeclaração que você disse que não precisava, mas eu achei muito legal.

• criar dados fictícios
• alterar coluna para nula
• alterar o tipo da coluna para int
• criar chave estrangeira
• definir um único valor fk
• verifique se o relacionamento está funcionando

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)