fix: Retrieve columns in compound indexes in correct order (#798)

Fixes #797

Co-authored-by: Knut Olav Løite <koloite@gmail.com>

Author: waltaskew

google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

@@ -1267,33 +1267,36 @@ def get_multi_indexes(
                i.table_schema,
                i.table_name,
                i.index_name,
-               (
-                   SELECT ARRAY_AGG(ic.column_name)
+               ARRAY(
+                   SELECT ic.column_name
                    FROM information_schema.index_columns ic
                    WHERE ic.index_name = i.index_name
                    AND ic.table_catalog = i.table_catalog
                    AND ic.table_schema = i.table_schema
                    AND ic.table_name = i.table_name
                    AND ic.column_ordering is not null
+                   ORDER BY ic.ordinal_position
                ) as columns,
                i.is_unique,
-               (
-                   SELECT ARRAY_AGG(ic.column_ordering)
+               ARRAY(
+                   SELECT ic.column_ordering
                    FROM information_schema.index_columns ic
                    WHERE ic.index_name = i.index_name
                    AND ic.table_catalog = i.table_catalog
                    AND ic.table_schema = i.table_schema
                    AND ic.table_name = i.table_name
                    AND ic.column_ordering is not null
+                   ORDER BY ic.ordinal_position
                ) as column_orderings,
-               (
-                   SELECT ARRAY_AGG(storing.column_name)
+               ARRAY(
+                   SELECT storing.column_name
                    FROM information_schema.index_columns storing
                    WHERE storing.index_name = i.index_name
                    AND storing.table_catalog = i.table_catalog
                    AND storing.table_schema = i.table_schema
                    AND storing.table_name = i.table_name
                    AND storing.column_ordering is null
+                   ORDER BY storing.ordinal_position
                ) as storing_columns,
             FROM information_schema.indexes as i
             JOIN information_schema.tables AS t

test/system/test_basics.py

@@ -106,7 +106,7 @@ def define_tables(cls, metadata):
             Column("a", String, primary_key=True),
             Column("b", String, primary_key=True),
         )
-        Table(
+        composite_fk = Table(
             "composite_fk",
             metadata,
             Column("my_a", String, primary_key=True),
@@ -118,6 +118,12 @@ def define_tables(cls, metadata):
                 name="composite_fk_composite_pk_a_b",
             ),
         )
+        Index(
+            "idx_composte_fk_all",
+            composite_fk.c.my_a,
+            composite_fk.c.my_b,
+            composite_fk.c.my_c,
+        )
 
     def test_hello_world(self, connection):
         greeting = connection.execute(text("select 'Hello World'"))
@@ -356,6 +362,28 @@ def test_composite_fk_lookups(self, connection):
             insp.get_multi_foreign_keys(filter_names=["composite_fk"]),
         )
 
+    def test_composite_index_lookups(self, connection):
+        """Ensures we introspect composite indexes."""
+
+        engine = connection.engine
+
+        insp = inspect(engine)
+        eq_(
+            {
+                (None, "composite_fk"): [
+                    {
+                        "name": "idx_composte_fk_all",
+                        "column_names": ["my_a", "my_b", "my_c"],
+                        "unique": False,
+                        "column_sorting": {"my_a": "asc", "my_b": "asc", "my_c": "asc"},
+                        "include_columns": [],
+                        "dialect_options": {},
+                    }
+                ]
+            },
+            insp.get_multi_indexes(filter_names=["composite_fk"]),
+        )
+
     def test_commit_timestamp(self, connection):
         """Ensures commit timestamps are set."""