Autogenerate loses variants on TypeDecorator column with SQLAlchemy 2.0

[mgorven]

Describe the bug
An autogenerated migration adding a column using a TypeDecorator type and a with_variant() override loses the variants when using SQLAlchemy 2.0.

Expected behavior
The autogenerated migration includes the variants.

To Reproduce

With the following model definition and then alembic revision --autogenerate.

from sqlalchemy import Column, String, types
from sqlalchemy.dialects import postgresql
from sqlalchemy.orm import declarative_base

Base = declarative_base()

class UUID(types.TypeDecorator):
    impl = String

class Test(Base):
    __tablename__ = "test"
    id = Column(UUID().with_variant(postgresql.UUID(), "postgresql"), primary_key=True)

Error

With SQLAlchemy 1.4.46 the following migration is generated:

    op.create_table('test',
    sa.Column('id', models.UUID().with_variant(postgresql.UUID(), 'postgresql'), nullable=False),
    sa.PrimaryKeyConstraint('id')
    )

Whereas with SQLALchemy 2.0.2 the with_variant() is missing:

    op.create_table('test',
    sa.Column('id', models.UUID(), nullable=False),
    sa.PrimaryKeyConstraint('id')
    )

Versions.

• OS: Linux
• Python: 3.10
• Alembic: 1.9.2
• SQLAlchemy: 2.0.2
• Database: SQLite
• DBAPI: sqlite

[zzzeek]

it's ....not really that simple. alembic knows how to render 2.0's variants, but this has to do with your datatype not being part of SQLAlchemy, and there seems to be a whole codepath that isn't taken for variants that are not within sqlalchemy's modules . it "works" in 1.4 because your type becomes Variant, which is in the normal module path.

but then...it is that simple? because this is is all it needs..

diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 41903d8..0a5c3dc 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -837,6 +837,8 @@ def _repr_type(
             prefix = _sqlalchemy_autogenerate_prefix(autogen_context)
             return "%s%r" % (prefix, type_)
     else:
+        if not _skip_variants and sqla_compat._type_has_variants(type_):
+            return _render_Variant_type(type_, autogen_context)
         prefix = _user_autogenerate_prefix(autogen_context, type_)
         return "%s%r" % (prefix, type_)
 

[sqla-tester]

Mike Bayer has proposed a fix for this issue in the main branch:

add variant render step for user-defined types https://gerrit.sqlalchemy.org/c/sqlalchemy/alembic/+/4431

[mgorven]

Thanks for the quick fix!