From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Tue, 16 Nov 2021 19:10:58 +0000 (-0500)
Subject: support 2.0 style variants
X-Git-Tag: rel_1_7_6~14
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=d18af2d6d1c49cc86b878d654b3404bb26614be6;p=thirdparty%2Fsqlalchemy%2Falembic.git

support 2.0 style variants

Implemented support for recognizing and rendering SQLAlchemy "variant"
types going forward into SQLAlchemy 2.0, where the architecture of
"variant" datatypes will be changing.

Specifically any TypeEngine can now have "variants" by looking
in the _variant_mapping dictionary collection associated with the
type directly.

Change-Id: I0b5ca887dce1dd77af3504dbe318701b31c7d574
---

diff --git a/alembic/autogenerate/render.py b/alembic/autogenerate/render.py
index 90d49e5f..b8226f7a 100644
--- a/alembic/autogenerate/render.py
+++ b/alembic/autogenerate/render.py
@@ -41,7 +41,6 @@ if TYPE_CHECKING:
     from sqlalchemy.sql.schema import UniqueConstraint
     from sqlalchemy.sql.sqltypes import ARRAY
     from sqlalchemy.sql.type_api import TypeEngine
-    from sqlalchemy.sql.type_api import Variant
 
     from alembic.autogenerate.api import AutogenContext
     from alembic.config import Config
@@ -812,7 +811,11 @@ def _get_identity_options(identity_options: "Identity") -> OrderedDict:
     return kwargs
 
 
-def _repr_type(type_: "TypeEngine", autogen_context: "AutogenContext") -> str:
+def _repr_type(
+    type_: "TypeEngine",
+    autogen_context: "AutogenContext",
+    _skip_variants: bool = False,
+) -> str:
     rendered = _user_defined_render("type", type_, autogen_context)
     if rendered is not False:
         return rendered
@@ -839,7 +842,7 @@ def _repr_type(type_: "TypeEngine", autogen_context: "AutogenContext") -> str:
     elif impl_rt:
         return impl_rt
     elif mod.startswith("sqlalchemy."):
-        if type(type_) is sqltypes.Variant:
+        if not _skip_variants and sqla_compat._type_has_variants(type_):
             return _render_Variant_type(type_, autogen_context)
         if "_render_%s_type" % type_.__visit_name__ in globals():
             fn = globals()["_render_%s_type" % type_.__visit_name__]
@@ -864,14 +867,15 @@ def _render_ARRAY_type(
 
 
 def _render_Variant_type(
-    type_: "Variant", autogen_context: "AutogenContext"
+    type_: "TypeEngine", autogen_context: "AutogenContext"
 ) -> str:
-    base = _repr_type(type_.impl, autogen_context)
+    base_type, variant_mapping = sqla_compat._get_variant_mapping(type_)
+    base = _repr_type(base_type, autogen_context, _skip_variants=True)
     assert base is not None and base is not False
-    for dialect in sorted(type_.mapping):
-        typ = type_.mapping[dialect]
+    for dialect in sorted(variant_mapping):
+        typ = variant_mapping[dialect]
         base += ".with_variant(%s, %r)" % (
-            _repr_type(typ, autogen_context),
+            _repr_type(typ, autogen_context, _skip_variants=True),
             dialect,
         )
     return base
diff --git a/alembic/util/sqla_compat.py b/alembic/util/sqla_compat.py
index 57a6a769..221e20e8 100644
--- a/alembic/util/sqla_compat.py
+++ b/alembic/util/sqla_compat.py
@@ -252,6 +252,24 @@ def _reflect_table(
         return inspector.reflecttable(table, None)
 
 
+if hasattr(sqltypes.TypeEngine, "_variant_mapping"):
+
+    def _type_has_variants(type_):
+        return bool(type_._variant_mapping)
+
+    def _get_variant_mapping(type_):
+        return type_, type_._variant_mapping
+
+
+else:
+
+    def _type_has_variants(type_):
+        return type(type_) is sqltypes.Variant
+
+    def _get_variant_mapping(type_):
+        return type_.impl, type_.mapping
+
+
 def _fk_spec(constraint):
     source_columns = [
         constraint.columns[key].name for key in constraint.column_keys
diff --git a/docs/build/unreleased/variant.rst b/docs/build/unreleased/variant.rst
new file mode 100644
index 00000000..71cf3c79
--- /dev/null
+++ b/docs/build/unreleased/variant.rst
@@ -0,0 +1,7 @@
+.. change::
+    :tags: bug, autogenerate
+
+    Implemented support for recognizing and rendering SQLAlchemy "variant"
+    types going forward into SQLAlchemy 2.0, where the architecture of
+    "variant" datatypes will be changing.
+