implement full copy for indexes in batch

Fixed issue in batch mode where CREATE INDEX would not use a new column
name in the case of a column rename.

Indexes were previously being moved to the new table without any
steps to rewrite columns.   We now vendor the copy approach
from table.to_metadata so that there's a new index expressed
in terms of the new columns.

Change-Id: Ied84232037aee0b2bf2094b3d3474013d7b41b34
Fixes: #1034

3 files changed, 29 insertions, 3 deletions

diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py
index 308bc2e..7c7de9f 100644
--- a/alembic/operations/batch.py
+++ b/alembic/operations/batch.py
@@ -24,8 +24,10 @@ from sqlalchemy.util import topological
 from ..util import exc
 from ..util.sqla_compat import _columns_for_constraint
 from ..util.sqla_compat import _copy
+from ..util.sqla_compat import _copy_expression
 from ..util.sqla_compat import _ensure_scope_for_ddl
 from ..util.sqla_compat import _fk_is_self_referential
+from ..util.sqla_compat import _idx_table_bound_expressions
 from ..util.sqla_compat import _insert_inline
 from ..util.sqla_compat import _is_type_bound
 from ..util.sqla_compat import _remove_column_from_collection
@@ -354,7 +356,25 @@ class ApplyBatchImpl:
     def _gather_indexes_from_both_tables(self) -> List["Index"]:
         assert self.new_table is not None
         idx: List[Index] = []
-        idx.extend(self.indexes.values())
+
+        for idx_existing in self.indexes.values():
+            # this is a lift-and-move from Table.to_metadata
+
+            if idx_existing._column_flag:  # type: ignore
+                continue
+
+            idx_copy = Index(
+                idx_existing.name,
+                unique=idx_existing.unique,
+                *[
+                    _copy_expression(expr, self.new_table)
+                    for expr in _idx_table_bound_expressions(idx_existing)
+                ],
+                _table=self.new_table,
+                **idx_existing.kwargs,
+            )
+            idx.append(idx_copy)
+
         for index in self.new_indexes.values():
             idx.append(
                 Index(
diff --git a/alembic/util/compat.py b/alembic/util/compat.py
index cabff6e..289aaa2 100644
--- a/alembic/util/compat.py
+++ b/alembic/util/compat.py
@@ -35,9 +35,9 @@ else:
 def importlib_metadata_get(group: str) -> Sequence[EntryPoint]:
     ep = importlib_metadata.entry_points()
     if hasattr(ep, "select"):
-        return ep.select(group=group)  # type:ignore[attr-defined]
+        return ep.select(group=group)  # type: ignore
     else:
-        return ep.get(group, ())
+        return ep.get(group, ())  # type: ignore
 
 
 def formatannotation_fwdref(annotation, base_module=None):
diff --git a/alembic/util/sqla_compat.py b/alembic/util/sqla_compat.py
index 21a9f7f..9c06e85 100644
--- a/alembic/util/sqla_compat.py
+++ b/alembic/util/sqla_compat.py
@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import contextlib
 import re
+from typing import Any
+from typing import Iterable
 from typing import Iterator
 from typing import Mapping
 from typing import Optional
@@ -158,6 +160,10 @@ def _get_connection_in_transaction(connection: Optional["Connection"]) -> bool:
         return in_transaction()
 
 
+def _idx_table_bound_expressions(idx: Index) -> Iterable[ColumnElement[Any]]:
+    return idx.expressions  # type: ignore
+
+
 def _copy(schema_item: _CE, **kw) -> _CE:
     if hasattr(schema_item, "_copy"):
         return schema_item._copy(**kw)  # type: ignore[union-attr]