fix(sqlalchemy-spanner): support spanner_storing in Alembic op.create_index() and compare_indexes()

packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

@@ -19,7 +19,7 @@
 from google.auth.credentials import AnonymousCredentials
 from google.cloud.spanner_v1 import Client, TransactionOptions
 from google.cloud.spanner_v1.data_types import JsonObject
-from sqlalchemy import ForeignKeyConstraint, PickleType, TypeDecorator, types
+from sqlalchemy import Column, ForeignKeyConstraint, PickleType, TypeDecorator, types
 from sqlalchemy.engine.base import Engine
 from sqlalchemy.engine.default import DefaultDialect, DefaultExecutionContext
 from sqlalchemy.event import listens_for
@@ -58,6 +58,7 @@
         format_server_default,
         format_type,
     )
+    from alembic.ddl.impl import ComparisonResult, DefaultImpl
 
     HAS_ALEMBIC_INSTALLED = True
 except ImportError:
@@ -740,12 +741,11 @@ def visit_create_index(
             options = index.dialect_options["spanner"]
             if "storing" in options:
                 storing = options["storing"]
-                storing_columns = [
-                    index.table.c[col] if isinstance(col, str) else col
-                    for col in storing
+                storing_col_names = [
+                    c if isinstance(c, str) else c.name for c in storing
                 ]
                 text += " STORING (%s)" % ", ".join(
-                    [self.preparer.quote(c.name) for c in storing_columns]
+                    [self.preparer.quote(name) for name in storing_col_names]
                 )
 
             interleave_in = options.get("interleave_in")
@@ -1955,3 +1955,36 @@ def _format_alter_column(
                 else f" DEFAULT {format_server_default(compiler, server_default)}"
             ),
         )
+
+    class SpannerImpl(DefaultImpl):
+        """Alembic DDL implementation for Google Cloud Spanner."""
+
+        __dialect__ = "spanner+spanner"
+
+        def compare_indexes(self, metadata_index, reflected_index) -> ComparisonResult:
+            result = super().compare_indexes(metadata_index, reflected_index)
+            if not result.is_equal:
+                return result
+
+            def _get_storing_columns(index):
+                storing = index.dialect_options.get("spanner", {}).get("storing")
+                if not storing:
+                    return []
+                return [col if isinstance(col, str) else col.name for col in storing]
+
+            m_storing = _get_storing_columns(metadata_index)
+            r_storing = _get_storing_columns(reflected_index)
+            if set(m_storing) != set(r_storing):
+                return ComparisonResult.Different(
+                    [f"spanner_storing columns {r_storing} to {m_storing}"]
+                )
+            return result
+
+        def create_index(self, index, **kw):
+            storing = index.dialect_options.get("spanner", {}).get("storing") or ()
+            if index.table is not None:
+                for col in storing:
+                    col_name = col if isinstance(col, str) else col.name
+                    if col_name not in index.table.c:
+                        index.table.append_column(Column(col_name, types.NullType))
+            super().create_index(index, **kw)

packages/sqlalchemy-spanner/tests/unit/test_alembic.py

@@ -13,11 +13,25 @@
 # limitations under the License.
 
 import importlib.util
+import io
 import sys
 from unittest import mock
 
+import pytest
 from alembic.ddl import base as ddl_base
-from sqlalchemy import String, TextClause, event
+from sqlalchemy import (
+    Column,
+    Index,
+    Integer,
+    MetaData,
+    String,
+    Table,
+    TextClause,
+    event,
+    exc,
+    schema,
+    text,
+)
 from sqlalchemy.testing import eq_
 from sqlalchemy.testing.plugin.plugin_base import fixtures
 
@@ -130,3 +144,147 @@ def test_dialect_import_without_alembic(self):
             event.remove(sqlalchemy_spanner.Pool, "reset", module.reset_connection)
             assert not module.HAS_ALEMBIC_INSTALLED
             assert not hasattr(module, "visit_column_nullable")
+            assert not hasattr(module, "SpannerImpl")
+
+    def test_visit_create_index_with_storing_columns(self):
+        compiler = sqlalchemy_spanner.SpannerDDLCompiler(
+            sqlalchemy_spanner.SpannerDialect(), None
+        )
+        t = Table("tbl", MetaData(), Column("c1", Integer), Column("c2", String(50)))
+        idx = Index("ix_test", t.c.c1, spanner_storing=["c2"])
+        ddl = compiler.visit_create_index(schema.CreateIndex(idx))
+        eq_(ddl, "CREATE INDEX ix_test ON tbl (c1) STORING (c2)")
+
+    def test_visit_create_index_with_storing_columns_missing_from_table(self):
+        compiler = sqlalchemy_spanner.SpannerDDLCompiler(
+            sqlalchemy_spanner.SpannerDialect(), None
+        )
+        # Mimics Alembic's op.create_index where table only has indexed columns
+        t = Table("tbl", MetaData(), Column("c1", Integer))
+        idx = Index("ix_test", t.c.c1, spanner_storing=["c2"])
+        ddl = compiler.visit_create_index(schema.CreateIndex(idx))
+        eq_(ddl, "CREATE INDEX ix_test ON tbl (c1) STORING (c2)")
+
+    def test_spanner_impl_create_index_appends_missing_columns(self):
+        from alembic.migration import MigrationContext
+
+        dialect = sqlalchemy_spanner.SpannerDialect()
+        impl = MigrationContext.configure(
+            dialect=dialect, opts={"as_sql": True, "output_buffer": io.StringIO()}
+        ).impl
+
+        t = Table("tbl", MetaData(), Column("c1", Integer))
+        idx = Index("ix_test", t.c.c1, spanner_storing=["c2", "c3"])
+        assert "c2" not in t.c
+        assert "c3" not in t.c
+
+        impl.create_index(idx)
+        assert "c2" in t.c
+        assert "c3" in t.c
+
+    def test_spanner_impl_compare_indexes(self):
+        from alembic.migration import MigrationContext
+
+        dialect = sqlalchemy_spanner.SpannerDialect()
+        impl = MigrationContext.configure(dialect=dialect).impl
+
+        t = Table(
+            "tbl",
+            MetaData(),
+            Column("c1", Integer),
+            Column("c2", String(50)),
+            Column("c3", String(50)),
+        )
+        ix_no_storing = Index("ix", t.c.c1)
+        ix_storing = Index("ix", t.c.c1, spanner_storing=["c2", "c3"])
+        ix_reordered = Index("ix", t.c.c1, spanner_storing=["c3", "c2"])
+        ix_different = Index("ix", t.c.c1, spanner_storing=["c2"])
+
+        # Identical or reordered storing columns compare equal
+        eq_(impl.compare_indexes(ix_no_storing, Index("ix", t.c.c1)).is_equal, True)
+        eq_(
+            impl.compare_indexes(
+                ix_storing, Index("ix", t.c.c1, spanner_storing=["c2", "c3"])
+            ).is_equal,
+            True,
+        )
+        eq_(impl.compare_indexes(ix_storing, ix_reordered).is_equal, True)
+
+        # Added, removed, or changed storing columns return is_different
+        eq_(impl.compare_indexes(ix_storing, ix_no_storing).is_different, True)
+        eq_(impl.compare_indexes(ix_no_storing, ix_storing).is_different, True)
+        eq_(impl.compare_indexes(ix_storing, ix_different).is_different, True)
+
+    def test_visit_create_index_with_column_objects_in_storing(self):
+        compiler = sqlalchemy_spanner.SpannerDDLCompiler(
+            sqlalchemy_spanner.SpannerDialect(), None
+        )
+        t = Table("tbl", MetaData(), Column("c1", Integer), Column("c2", String(50)))
+        idx = Index("ix_test", t.c.c1, spanner_storing=[t.c.c2])
+        ddl = compiler.visit_create_index(schema.CreateIndex(idx))
+        eq_(ddl, "CREATE INDEX ix_test ON tbl (c1) STORING (c2)")
+
+    def test_spanner_impl_compare_indexes_super_different(self):
+        from alembic.migration import MigrationContext
+
+        dialect = sqlalchemy_spanner.SpannerDialect()
+        impl = MigrationContext.configure(dialect=dialect).impl
+
+        t = Table(
+            "tbl",
+            MetaData(),
+            Column("c1", Integer),
+            Column("c2", Integer),
+        )
+        # Key column differs
+        ix1 = Index("ix", t.c.c1)
+        ix2 = Index("ix", t.c.c2)
+        eq_(impl.compare_indexes(ix1, ix2).is_different, True)
+
+        # Uniqueness differs
+        ix_unique = Index("ix", t.c.c1, unique=True)
+        ix_non_unique = Index("ix", t.c.c1, unique=False)
+        eq_(impl.compare_indexes(ix_unique, ix_non_unique).is_different, True)
+
+    def test_spanner_impl_compare_indexes_with_column_objects(self):
+        from alembic.migration import MigrationContext
+
+        dialect = sqlalchemy_spanner.SpannerDialect()
+        impl = MigrationContext.configure(dialect=dialect).impl
+
+        t = Table(
+            "tbl",
+            MetaData(),
+            Column("c1", Integer),
+            Column("c2", Integer),
+        )
+        ix1 = Index("ix", t.c.c1, spanner_storing=[t.c.c2])
+        ix2 = Index("ix", t.c.c1, spanner_storing=["c2"])
+        eq_(impl.compare_indexes(ix1, ix2).is_equal, True)
+
+    def test_spanner_impl_create_index_edge_cases(self):
+        from alembic.migration import MigrationContext
+
+        dialect = sqlalchemy_spanner.SpannerDialect()
+        impl = MigrationContext.configure(
+            dialect=dialect, opts={"as_sql": True, "output_buffer": io.StringIO()}
+        ).impl
+
+        # 1. col already exists in table.c
+        t1 = Table("tbl1", MetaData(), Column("c1", Integer), Column("c2", Integer))
+        idx1 = Index("ix_test1", t1.c.c1, spanner_storing=["c2"])
+        impl.create_index(idx1)
+        eq_(len([c for c in t1.c if c.name == "c2"]), 1)
+
+        # 2. col is a Column object not in table.c
+        t2 = Table("tbl2", MetaData(), Column("c1", Integer))
+        c_ext = Column("c_ext", Integer)
+        idx2 = Index("ix_test2", t2.c.c1, spanner_storing=[c_ext])
+        impl.create_index(idx2)
+        assert "c_ext" in t2.c
+
+        # 3. index.table is None
+        idx3 = Index("ix_test3", text("c1"), spanner_storing=["c2"])
+        idx3.table = None
+        with pytest.raises(exc.CompileError):
+            impl.create_index(idx3)