[spanner-sqlalchemy]: support spanner_storing in Alembic op.create_index() and autogenerate compare_indexes() #18835

if "storing" in options:
    storing = options["storing"]
    storing_columns = [
        index.table.c[col] if isinstance(col, str) else col
        for col in storing
    ]
    text += " STORING (%s)" % ", ".join(
        [self.preparer.quote(c.name) for c in storing_columns]
    )

import io
import alembic.operations
from alembic.ddl.impl import DefaultImpl
from alembic.migration import MigrationContext
from google.cloud.sqlalchemy_spanner.sqlalchemy_spanner import SpannerDialect
import sqlalchemy as sa


class SpannerDefaultImpl(DefaultImpl):
    __dialect__ = "spanner+spanner"


# 1. Reproduce KeyError in op.create_index with spanner_storing=["c2"]
buf = io.StringIO()
ctx = MigrationContext.configure(
    dialect=SpannerDialect(),
    opts={"as_sql": True, "output_buffer": buf},
)
op = alembic.operations.Operations(ctx)

try:
    op.create_index("ix_tbl_c1", "tbl", ["c1"], spanner_storing=["c2"])
except KeyError as e:
    print(f"Issue 1 reproduced: KeyError: {e!r}")

# 2. Reproduce compare_indexes ignoring spanner_storing changes
t1 = sa.Table("tbl", sa.MetaData(), sa.Column("c1", sa.Integer), sa.Column("c2", sa.String))
idx_meta = sa.Index("ix_tbl_c1", t1.c.c1, spanner_storing=["c2"])

t2 = sa.Table("tbl", sa.MetaData(), sa.Column("c1", sa.Integer), sa.Column("c2", sa.String))
idx_refl = sa.Index("ix_tbl_c1", t2.c.c1)

print("Issue 2 compare_indexes result:", ctx.impl.compare_indexes(idx_meta, idx_refl))

"""MCVE reproducing two issues in google-cloud-sqlalchemy-spanner and alembic:

Issue 1: SpannerDDLCompiler.visit_create_index fails with KeyError when op.create_index
         is called with spanner_storing=[<str>, ...] because Alembic creates an ad-hoc Table
         containing only index key columns, so index.table.c[col] fails.
         Location: google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py:1203

Issue 2: Alembic's DefaultImpl (used by SpannerDialect) ignores spanner_storing in
         compare_indexes, failing to detect when storing columns are added, removed, or changed.
         google-cloud-sqlalchemy-spanner currently does not provide an Alembic Impl subclass
         to compare dialect-specific index options.
"""

from __future__ import annotations

import io
import sys
import alembic.operations
from alembic.ddl.impl import DefaultImpl
from alembic.migration import MigrationContext
from google.cloud.sqlalchemy_spanner.sqlalchemy_spanner import SpannerDialect
import sqlalchemy as sa


# Alembic requires a DefaultImpl subclass registered for the dialect name.
# google-cloud-sqlalchemy-spanner does not provide one out of the box.
class SpannerDefaultImpl(DefaultImpl):
    __dialect__ = "spanner+spanner"


def reproduce_issue_1() -> None:
    """Demonstrate KeyError when op.create_index is called with spanner_storing."""
    print("=" * 70)
    print("Issue 1: op.create_index(..., spanner_storing=['c2']) raises KeyError")
    print("=" * 70)

    buf = io.StringIO()
    ctx = MigrationContext.configure(
        dialect=SpannerDialect(),
        opts={"as_sql": True, "output_buffer": buf},
    )
    op = alembic.operations.Operations(ctx)

    try:
        # In an Alembic migration downgrade() or upgrade(), op.create_index creates an Index
        # on an ad-hoc Table containing only the indexed column ['c1']. Column 'c2' is not in table.c.
        op.create_index("ix_tbl_c1", "tbl", ["c1"], spanner_storing=["c2"])
        print("UNEXPECTED: op.create_index succeeded!")
    except KeyError as e:
        print(f"REPRODUCED: KeyError: {e!r}")
        print("Root cause:")
        print("  In google.cloud.sqlalchemy_spanner.SpannerDDLCompiler.visit_create_index:")
        print("    storing_columns = [")
        print("        index.table.c[col] if isinstance(col, str) else col")
        print("        for col in storing")
        print("    ]")
        print("  index.table.c['c2'] fails because 'c2' is only in STORING, not an indexed column.")


def reproduce_issue_2() -> None:
    """Demonstrate compare_indexes ignoring spanner_storing changes."""
    print("\n" + "=" * 70)
    print("Issue 2: DefaultImpl.compare_indexes ignores spanner_storing changes")
    print("=" * 70)

    dialect = SpannerDialect()
    buf = io.StringIO()
    ctx = MigrationContext.configure(
        dialect=dialect,
        opts={"as_sql": True, "output_buffer": buf},
    )
    impl = ctx.impl

    # Metadata index has storing=['c2']
    t1 = sa.Table("tbl", sa.MetaData(), sa.Column("c1", sa.Integer), sa.Column("c2", sa.String))
    idx_meta = sa.Index("ix_tbl_c1", t1.c.c1, spanner_storing=["c2"])

    # Reflected database index has NO storing columns
    t2 = sa.Table("tbl", sa.MetaData(), sa.Column("c1", sa.Integer), sa.Column("c2", sa.String))
    idx_refl = sa.Index("ix_tbl_c1", t2.c.c1)

    result = impl.compare_indexes(idx_meta, idx_refl)
    print(f"Result of compare_indexes: {result!r}")
    if getattr(result, "status", None) == "equal" or result == 0:
        print("REPRODUCED: Indexes are reported EQUAL despite different spanner_storing!")
        print("Root cause:")
        print("  DefaultImpl.compare_indexes only compares unique constraint and indexed columns.")
        print("  Spanner-specific storing columns require an override in a SpannerImpl class.")


if __name__ == "__main__":
    reproduce_issue_1()
    reproduce_issue_2()