Improved support for expression indexes

Added support for autogenerate comparison of indexes on PostgreSQL which
include SQL expressions; the previous warning that such indexes were
skipped is now removed. This functionality requires SQLAlchemy 2.0.
For older SQLAlchemy versions, these indexes are still skipped.

Fixed issue where indexes on SQLite which include SQL expressions would not
compare against themselves correctly, generating false positives.
SQLAlchemy as of version 2 has no support for reflecting expression based
indexes on SQLite; so for now, the behavior is that SQLite expression-based
indexes are ignored for autogenerate compare, in the same way that
PostgreSQL expression-based indexes were ignored for the time that
SQLAlchemy did not support reflection of such indexes (which is now
supported in SQLAlchemy 2.0 as well as this release of Alembic).

Fixed issue in index detection where autogenerate change detection would
consider indexes with the same columns but with different order as equal,
while in general they are not equivalent in how a database will use them.

Fixes: #1165
Fixes: #1166
Change-Id: I226408eed855b923172e5df0bdab005ed2cc9f53

4 files changed, 374 insertions, 62 deletions

diff --git a/tests/requirements.py b/tests/requirements.py
index c774e67..1a100dd 100644
--- a/tests/requirements.py
+++ b/tests/requirements.py
@@ -410,3 +410,17 @@ class DefaultRequirements(SuiteRequirements):
         )
 
         return imports + version + sqlalchemy
+
+    @property
+    def reflect_indexes_with_expressions(self):
+        sqlalchemy = exclusions.only_if(
+            lambda _: sqla_compat.sqla_2, "sqlalchemy 2 is required"
+        )
+
+        postgresql = exclusions.only_on(["postgresql"])
+
+        return sqlalchemy + postgresql
+
+    @property
+    def indexes_with_expressions(self):
+        return exclusions.only_on(["postgresql", "sqlite>=3.9.0"])
diff --git a/tests/test_autogen_indexes.py b/tests/test_autogen_indexes.py
index 68a6bd6..6faba04 100644
--- a/tests/test_autogen_indexes.py
+++ b/tests/test_autogen_indexes.py
@@ -1,6 +1,11 @@
+from contextlib import nullcontext
+import itertools
+
 from sqlalchemy import Column
+from sqlalchemy import Float
 from sqlalchemy import ForeignKey
 from sqlalchemy import ForeignKeyConstraint
+from sqlalchemy import func
 from sqlalchemy import Index
 from sqlalchemy import Integer
 from sqlalchemy import MetaData
@@ -9,16 +14,20 @@ from sqlalchemy import PrimaryKeyConstraint
 from sqlalchemy import String
 from sqlalchemy import Table
 from sqlalchemy import UniqueConstraint
+from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION
 from sqlalchemy.sql.expression import column
 from sqlalchemy.sql.expression import desc
 
+from alembic import testing
 from alembic.testing import combinations
 from alembic.testing import config
 from alembic.testing import eq_
 from alembic.testing import exclusions
+from alembic.testing import resolve_lambda
 from alembic.testing import schemacompare
 from alembic.testing import TestBase
 from alembic.testing import util
+from alembic.testing.assertions import expect_warnings
 from alembic.testing.env import staging_env
 from alembic.testing.suite._autogen_fixtures import AutogenFixtureTest
 from alembic.util import sqla_compat
@@ -1175,6 +1184,287 @@ class AutogenerateIndexTest(AutogenFixtureTest, TestBase):
         diffs = self._fixture(m1, m2)
         eq_(diffs, [])
 
+    def test_column_order_changed(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        old = Index("SomeIndex", "x", "y")
+        Table(
+            "order_change",
+            m1,
+            Column("id", Integer, primary_key=True),
+            Column("x", Integer),
+            Column("y", Integer),
+            old,
+        )
+
+        new = Index("SomeIndex", "y", "x")
+        Table(
+            "order_change",
+            m2,
+            Column("id", Integer, primary_key=True),
+            Column("x", Integer),
+            Column("y", Integer),
+            new,
+        )
+        diffs = self._fixture(m1, m2)
+        eq_(
+            diffs,
+            [
+                ("remove_index", schemacompare.CompareIndex(old)),
+                ("add_index", schemacompare.CompareIndex(new)),
+            ],
+        )
+
+
+class AutogenerateExpressionIndexTest(AutogenFixtureTest, TestBase):
+    """tests involving indexes with expression"""
+
+    __requires__ = ("indexes_with_expressions",)
+
+    __backend__ = True
+
+    @property
+    def has_reflection(self):
+        return config.requirements.reflect_indexes_with_expressions.enabled
+
+    def test_expression_indexes_add(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        Table(
+            "exp_index",
+            m1,
+            Column("id", Integer, primary_key=True),
+            Column("x", Integer),
+            Column("y", Integer),
+        )
+
+        idx = Index("SomeIndex", "y", func.lower(column("x")))  # noqa
+        Table(
+            "exp_index",
+            m2,
+            Column("id", Integer, primary_key=True),
+            Column("x", Integer),
+            Column("y", Integer),
+            idx,
+        )
+
+        if self.has_reflection:
+            diffs = self._fixture(m1, m2)
+            eq_(diffs, [("add_index", schemacompare.CompareIndex(idx))])
+        else:
+            with expect_warnings(
+                r"autogenerate skipping metadata-specified expression-based "
+                r"index 'SomeIndex'; dialect '.*' under SQLAlchemy .* "
+                r"can't reflect these "
+                r"indexes so they can't be compared",
+            ):
+                diffs = self._fixture(m1, m2)
+            eq_(diffs, [])
+
+    def _lots_of_indexes(flatten: bool = False):
+        diff_pairs = [
+            (
+                lambda t: Index("SomeIndex", "y", func.lower(t.c.x)),
+                lambda t: Index("SomeIndex", func.lower(t.c.x)),
+            ),
+            (
+                lambda CapT: Index("SomeIndex", "y", func.lower(CapT.c.XCol)),
+                lambda CapT: Index("SomeIndex", func.lower(CapT.c.XCol)),
+            ),
+            (
+                lambda t: Index(
+                    "SomeIndex", "y", func.lower(column("x")), _table=t
+                ),
+                lambda t: Index(
+                    "SomeIndex", func.lower(column("x")), _table=t
+                ),
+            ),
+            (
+                lambda t: Index("SomeIndex", t.c.y, func.lower(t.c.x)),
+                lambda t: Index("SomeIndex", func.lower(t.c.x), t.c.y),
+            ),
+            (
+                lambda t: Index("SomeIndex", t.c.y, func.lower(t.c.x)),
+                lambda t: Index("SomeIndex", t.c.y, func.lower(t.c.q)),
+            ),
+            (
+                lambda t: Index("SomeIndex", t.c.z, func.lower(t.c.x)),
+                lambda t: Index("SomeIndex", t.c.y, func.lower(t.c.x)),
+            ),
+            (
+                lambda t: Index("SomeIndex", func.lower(t.c.x)),
+                lambda t: Index("SomeIndex", t.c.y, func.lower(t.c.x)),
+            ),
+            (
+                lambda t: Index("SomeIndex", t.c.y, func.upper(t.c.x)),
+                lambda t: Index("SomeIndex", t.c.y, func.lower(t.c.x)),
+            ),
+            (
+                lambda t: Index("SomeIndex", t.c.y, t.c.ff + 1),
+                lambda t: Index("SomeIndex", t.c.y, t.c.ff + 3),
+            ),
+            (
+                lambda t: Index("SomeIndex", t.c.y, func.ceil(t.c.ff)),
+                lambda t: Index("SomeIndex", t.c.y, func.floor(t.c.ff)),
+            ),
+            (
+                lambda t: Index("SomeIndex", t.c.y, func.lower(t.c.x)),
+                lambda t: Index("SomeIndex", t.c.y, func.lower(t.c.x + t.c.q)),
+            ),
+            (
+                lambda t: Index("SomeIndex", t.c.y, t.c.z + 3),
+                lambda t: Index("SomeIndex", t.c.y, t.c.z * 3),
+            ),
+            (
+                lambda t: Index("SomeIndex", func.lower(t.c.x), t.c.q + "42"),
+                lambda t: Index("SomeIndex", func.lower(t.c.q), t.c.x + "42"),
+            ),
+            (
+                lambda t: Index("SomeIndex", func.lower(t.c.x), t.c.z + 42),
+                lambda t: Index("SomeIndex", t.c.z + 42, func.lower(t.c.q)),
+            ),
+            (
+                lambda t: Index("SomeIndex", t.c.ff + 42),
+                lambda t: Index("SomeIndex", 42 + t.c.ff),
+            ),
+        ]
+        if flatten:
+            return list(itertools.chain.from_iterable(diff_pairs))
+        else:
+            return diff_pairs
+
+    @testing.fixture
+    def index_changed_tables(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        t_old = Table(
+            "exp_index",
+            m1,
+            Column("id", Integer, primary_key=True),
+            Column("x", String(100)),
+            Column("y", String(100)),
+            Column("q", String(100)),
+            Column("z", Integer),
+            Column("ff", Float().with_variant(DOUBLE_PRECISION, "postgresql")),
+        )
+
+        t_new = Table(
+            "exp_index",
+            m2,
+            Column("id", Integer, primary_key=True),
+            Column("x", String(100)),
+            Column("y", String(100)),
+            Column("q", String(100)),
+            Column("z", Integer),
+            Column("ff", Float().with_variant(DOUBLE_PRECISION, "postgresql")),
+        )
+
+        CapT_old = Table(
+            "CapT table",
+            m1,
+            Column("id", Integer, primary_key=True),
+            Column("XCol", String(100)),
+            Column("y", String(100)),
+        )
+
+        CapT_new = Table(
+            "CapT table",
+            m2,
+            Column("id", Integer, primary_key=True),
+            Column("XCol", String(100)),
+            Column("y", String(100)),
+        )
+
+        return (
+            m1,
+            m2,
+            {"t": t_old, "CapT": CapT_old},
+            {"t": t_new, "CapT": CapT_new},
+        )
+
+    @combinations(*_lots_of_indexes(), argnames="old_fn, new_fn")
+    def test_expression_indexes_changed(
+        self, index_changed_tables, old_fn, new_fn
+    ):
+        m1, m2, old_fixture_tables, new_fixture_tables = index_changed_tables
+
+        old, new = resolve_lambda(
+            old_fn, **old_fixture_tables
+        ), resolve_lambda(new_fn, **new_fixture_tables)
+
+        if self.has_reflection:
+            diffs = self._fixture(m1, m2)
+            eq_(
+                diffs,
+                [
+                    ("remove_index", schemacompare.CompareIndex(old, True)),
+                    ("add_index", schemacompare.CompareIndex(new)),
+                ],
+            )
+        else:
+            with expect_warnings(
+                r"Skipped unsupported reflection of expression-based index "
+                r"SomeIndex",
+                r"autogenerate skipping metadata-specified expression-based "
+                r"index 'SomeIndex'; dialect '.*' under SQLAlchemy .* "
+                r"can't reflect these "
+                r"indexes so they can't be compared",
+            ):
+                diffs = self._fixture(m1, m2)
+            eq_(diffs, [])
+
+    @combinations(*_lots_of_indexes(flatten=True), argnames="fn")
+    def test_expression_indexes_no_change(self, index_changed_tables, fn):
+        m1, m2, old_fixture_tables, new_fixture_tables = index_changed_tables
+
+        resolve_lambda(fn, **old_fixture_tables)
+        resolve_lambda(fn, **new_fixture_tables)
+
+        if self.has_reflection:
+            ctx = nullcontext()
+        else:
+            ctx = expect_warnings()
+
+        with ctx:
+            diffs = self._fixture(m1, m2)
+        eq_(diffs, [])
+
+    def test_expression_indexes_remove(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        idx = Index("SomeIndex", "y", func.lower(column("x")))
+        Table(
+            "exp_index",
+            m1,
+            Column("id", Integer, primary_key=True),
+            Column("x", String(100)),
+            Column("y", Integer),
+            idx,
+        )
+
+        Table(
+            "exp_index",
+            m2,
+            Column("id", Integer, primary_key=True),
+            Column("x", String(100)),
+            Column("y", Integer),
+        )
+
+        if self.has_reflection:
+            diffs = self._fixture(m1, m2)
+            eq_(
+                diffs,
+                [("remove_index", schemacompare.CompareIndex(idx, True))],
+            )
+        else:
+            with expect_warnings():
+                diffs = self._fixture(m1, m2)
+            eq_(diffs, [])
+
 
 class NoUqReflectionIndexTest(NoUqReflection, AutogenerateUniqueIndexTest):
     __only_on__ = "sqlite"
diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py
index 6a67e0b..c9860c2 100644
--- a/tests/test_postgresql.py
+++ b/tests/test_postgresql.py
@@ -39,7 +39,6 @@ from alembic.migration import MigrationContext
 from alembic.operations import ops
 from alembic.script import ScriptDirectory
 from alembic.testing import assert_raises_message
-from alembic.testing import assertions
 from alembic.testing import combinations
 from alembic.testing import config
 from alembic.testing import eq_
@@ -1305,64 +1304,3 @@ class PGUniqueIndexAutogenerateTest(AutogenFixtureTest, TestBase):
         eq_(diffs[0][0], "remove_constraint")
         eq_(diffs[0][1].name, "uq_name")
         eq_(len(diffs), 1)
-
-    def _functional_index_warn(self):
-        return (r"Skip.*refl",)
-
-    def test_functional_ix_one(self):
-        m1 = MetaData()
-        m2 = MetaData()
-
-        t1 = Table(
-            "foo",
-            m1,
-            Column("id", Integer, primary_key=True),
-            Column("email", String(50)),
-        )
-        Index("email_idx", func.lower(t1.c.email), unique=True)
-
-        t2 = Table(
-            "foo",
-            m2,
-            Column("id", Integer, primary_key=True),
-            Column("email", String(50)),
-        )
-        Index("email_idx", func.lower(t2.c.email), unique=True)
-
-        with assertions.expect_warnings(*self._functional_index_warn()):
-            diffs = self._fixture(m1, m2)
-        eq_(diffs, [])
-
-    def test_functional_ix_two(self):
-        m1 = MetaData()
-        m2 = MetaData()
-
-        t1 = Table(
-            "foo",
-            m1,
-            Column("id", Integer, primary_key=True),
-            Column("email", String(50)),
-            Column("name", String(50)),
-        )
-        Index(
-            "email_idx",
-            func.coalesce(t1.c.email, t1.c.name).desc(),
-            unique=True,
-        )
-
-        t2 = Table(
-            "foo",
-            m2,
-            Column("id", Integer, primary_key=True),
-            Column("email", String(50)),
-            Column("name", String(50)),
-        )
-        Index(
-            "email_idx",
-            func.coalesce(t2.c.email, t2.c.name).desc(),
-            unique=True,
-        )
-
-        with assertions.expect_warnings(*self._functional_index_warn()):
-            diffs = self._fixture(m1, m2)
-        eq_(diffs, [])
diff --git a/tests/test_sqlite.py b/tests/test_sqlite.py
index b8edeaa..4c60ce4 100644
--- a/tests/test_sqlite.py
+++ b/tests/test_sqlite.py
@@ -3,6 +3,7 @@ from sqlalchemy import Column
 from sqlalchemy import DateTime
 from sqlalchemy import Float
 from sqlalchemy import func
+from sqlalchemy import Index
 from sqlalchemy import inspect
 from sqlalchemy import Integer
 from sqlalchemy import MetaData
@@ -21,10 +22,12 @@ from alembic.testing import assert_raises_message
 from alembic.testing import config
 from alembic.testing import eq_
 from alembic.testing import eq_ignore_whitespace
+from alembic.testing.assertions import expect_warnings
 from alembic.testing.env import clear_staging_env
 from alembic.testing.env import staging_env
 from alembic.testing.fixtures import op_fixture
 from alembic.testing.fixtures import TestBase
+from alembic.testing.suite._autogen_fixtures import AutogenFixtureTest
 
 
 class SQLiteTest(TestBase):
@@ -274,3 +277,70 @@ class SQLiteAutogenRenderTest(TestBase):
             "sa.Column('int_value', sa.Integer(), "
             "nullable=True, sqlite_on_conflict_not_null='FAIL')",
         )
+
+
+class SQLiteAutogenIndexTest(AutogenFixtureTest, TestBase):
+    __requires__ = ("indexes_with_expressions",)
+    __only_on__ = "sqlite"
+    __backend__ = True
+
+    def _functional_index_warn(self):
+        return (r"Skip.*refl",)
+
+    def test_functional_ix_one(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        t1 = Table(
+            "foo",
+            m1,
+            Column("id", Integer, primary_key=True),
+            Column("email", String(50)),
+        )
+        Index("email_idx", func.lower(t1.c.email), unique=True)
+
+        t2 = Table(
+            "foo",
+            m2,
+            Column("id", Integer, primary_key=True),
+            Column("email", String(50)),
+        )
+        Index("email_idx", func.lower(t2.c.email), unique=True)
+
+        with expect_warnings(*self._functional_index_warn()):
+            diffs = self._fixture(m1, m2)
+        eq_(diffs, [])
+
+    def test_functional_ix_two(self):
+        m1 = MetaData()
+        m2 = MetaData()
+
+        t1 = Table(
+            "foo",
+            m1,
+            Column("id", Integer, primary_key=True),
+            Column("email", String(50)),
+            Column("name", String(50)),
+        )
+        Index(
+            "email_idx",
+            func.coalesce(t1.c.email, t1.c.name).desc(),
+            unique=True,
+        )
+
+        t2 = Table(
+            "foo",
+            m2,
+            Column("id", Integer, primary_key=True),
+            Column("email", String(50)),
+            Column("name", String(50)),
+        )
+        Index(
+            "email_idx",
+            func.coalesce(t2.c.email, t2.c.name).desc(),
+            unique=True,
+        )
+
+        with expect_warnings(*self._functional_index_warn()):
+            diffs = self._fixture(m1, m2)
+        eq_(diffs, [])