fix: introspect constraints, keeping their order (#289)

Ilya Gurov · web-flow · commit 05cc7a17bec0 · 2022-11-23T22:52:53.000-08:00
diff --git a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py
@@ -794,7 +794,13 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kw):
     ctu.table_name,
     ctu.table_schema,
     ARRAY_AGG(DISTINCT ccu.column_name),
-    ARRAY_AGG(kcu.column_name)
+    ARRAY_AGG(
+        DISTINCT CONCAT(
+            CAST(kcu.ordinal_position AS STRING),
+            '_____',
+            kcu.column_name
+        )
+    )
 FROM information_schema.table_constraints AS tc
 JOIN information_schema.constraint_column_usage AS ccu
     ON ccu.constraint_name = tc.constraint_name
@@ -815,6 +821,21 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kw):
             rows = snap.execute_sql(sql)
 
             for row in rows:
+                # Due to Spanner limitations, arrays order is not guaranteed during
+                # aggregation. Still, for constraints it's vital to keep the order
+                # of the referred columns, otherwise SQLAlchemy and Alembic may start
+                # to occasionally drop and recreate constraints. To avoid this, the
+                # method uses prefixes with the `key_column_usage.ordinal_position`
+                # values to ensure the columns are aggregated into an array in the
+                # correct order. Prefixes are only used under the hood. For more details
+                # see the issue:
+                # https://github.com/googleapis/python-spanner-sqlalchemy/issues/271
+                #
+                # The solution seem a bit clumsy, and should be improved as soon as a
+                # better approach found.
+                for index, value in enumerate(sorted(row[4])):
+                    row[4][index] = value.split("_____")[1]
+
                 keys.append(
                     {
                         "name": row[0],
@@ -824,6 +845,7 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kw):
                         "constrained_columns": row[4],
                     }
                 )
+
         return keys
 
     @engine_to_connection
