From 9c5c1de6c0e65002e29ffdd2ee58f7a6fdec289c Mon Sep 17 00:00:00 2001 From: Kaushal Prajapati Date: Mon, 25 Sep 2023 16:26:53 +0530 Subject: [PATCH] preserve the sequence order of foreign keys when dealing with composite keys. --- src/snowflake/sqlalchemy/snowdialect.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/snowflake/sqlalchemy/snowdialect.py b/src/snowflake/sqlalchemy/snowdialect.py index 350027f4..949e556f 100644 --- a/src/snowflake/sqlalchemy/snowdialect.py +++ b/src/snowflake/sqlalchemy/snowdialect.py @@ -413,7 +413,7 @@ def _get_schema_foreign_keys(self, connection, schema, **kw): referred_schema = self.normalize_name(row._mapping["pk_schema_name"]) foreign_key_map[name] = { "constrained_columns": [ - self.normalize_name(row._mapping["fk_column_name"]) + (row._mapping["key_sequence"], self.normalize_name(row._mapping["fk_column_name"])) ], # referred schema should be None in context where it doesn't need to be specified # https://docs.sqlalchemy.org/en/14/core/reflection.html#reflection-schema-qualified-interaction @@ -427,7 +427,7 @@ def _get_schema_foreign_keys(self, connection, schema, **kw): row._mapping["pk_table_name"] ), "referred_columns": [ - self.normalize_name(row._mapping["pk_column_name"]) + (row._mapping["key_sequence"], self.normalize_name(row._mapping["pk_column_name"])) ], "name": name, "table_name": self.normalize_name(row._mapping["fk_table_name"]), @@ -444,15 +444,18 @@ def _get_schema_foreign_keys(self, connection, schema, **kw): foreign_key_map[name]["options"] = options else: foreign_key_map[name]["constrained_columns"].append( - self.normalize_name(row._mapping["fk_column_name"]) + (row._mapping["key_sequence"], self.normalize_name(row._mapping["fk_column_name"])) ) foreign_key_map[name]["referred_columns"].append( - self.normalize_name(row._mapping["pk_column_name"]) + (row._mapping["key_sequence"], self.normalize_name(row._mapping["pk_column_name"])) ) ans = {} for _, v in foreign_key_map.items(): + v["constrained_columns"] = [col for i, col in sorted(v["constrained_columns"])] + v["referred_columns"] = [col for i, col in sorted(v["referred_columns"])] + if v["table_name"] not in ans: ans[v["table_name"]] = [] ans[v["table_name"]].append(