From: fulpm <8397318+fulpm@users.noreply.github.com>
Date: Tue, 20 Oct 2020 21:54:50 +0000 (-0400)
Subject: Correct reflection for composite primary keys
X-Git-Tag: rel_1_4_0b1~16^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=7937a409f9e685341b6a5b8385491befe0aaa4d4;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

Correct reflection for composite primary keys

Fixes: #5661

### Description
Fixes reflection of composite primary keys to maintain the correct column order in the MSSQL
and SQLite dialects.

Closes: #5662
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/5662
Pull-request-sha: b568dec7070b4f3ee46a528bdf16fb237baade2a

Change-Id: I452b23cbf7f389c4a0a34cffce5c32498efe37d2
---

diff --git a/doc/build/changelog/unreleased_13/5661.rst b/doc/build/changelog/unreleased_13/5661.rst
new file mode 100644
index 0000000000..042b6a4932
--- /dev/null
+++ b/doc/build/changelog/unreleased_13/5661.rst
@@ -0,0 +1,6 @@
+.. change::
+    :tags: mssql, sqlite, reflection
+    :tickets: 5661
+
+    Fixed issue with composite primary key columns not being reported
+    in the correct order. Patch courtesy @fulpm.
diff --git a/lib/sqlalchemy/dialects/mssql/base.py b/lib/sqlalchemy/dialects/mssql/base.py
index c8f2b4ca3b..a224c00bbd 100644
--- a/lib/sqlalchemy/dialects/mssql/base.py
+++ b/lib/sqlalchemy/dialects/mssql/base.py
@@ -3168,15 +3168,19 @@ class MSDialect(default.DefaultDialect):
         C = ischema.key_constraints.alias("C")
 
         # Primary key constraints
-        s = sql.select(
-            C.c.column_name, TC.c.constraint_type, C.c.constraint_name
-        ).where(
-            sql.and_(
-                TC.c.constraint_name == C.c.constraint_name,
-                TC.c.table_schema == C.c.table_schema,
-                C.c.table_name == tablename,
-                C.c.table_schema == owner,
-            ),
+        s = (
+            sql.select(
+                C.c.column_name, TC.c.constraint_type, C.c.constraint_name
+            )
+            .where(
+                sql.and_(
+                    TC.c.constraint_name == C.c.constraint_name,
+                    TC.c.table_schema == C.c.table_schema,
+                    C.c.table_name == tablename,
+                    C.c.table_schema == owner,
+                ),
+            )
+            .order_by(TC.c.constraint_name, C.c.ordinal_position)
         )
         c = connection.execution_options(future_result=True).execute(s)
         constraint_name = None
diff --git a/lib/sqlalchemy/dialects/sqlite/base.py b/lib/sqlalchemy/dialects/sqlite/base.py
index 8ef35514ab..8a4fbe8e55 100644
--- a/lib/sqlalchemy/dialects/sqlite/base.py
+++ b/lib/sqlalchemy/dialects/sqlite/base.py
@@ -1853,6 +1853,7 @@ class SQLiteDialect(default.DefaultDialect):
             constraint_name = result.group(1) if result else None
 
         cols = self.get_columns(connection, table_name, schema, **kw)
+        cols.sort(key=lambda col: col.get("primary_key"))
         pkeys = []
         for col in cols:
             if col["primary_key"]:
diff --git a/lib/sqlalchemy/testing/suite/test_reflection.py b/lib/sqlalchemy/testing/suite/test_reflection.py
index f8f93b5635..49aafa879c 100644
--- a/lib/sqlalchemy/testing/suite/test_reflection.py
+++ b/lib/sqlalchemy/testing/suite/test_reflection.py
@@ -1581,6 +1581,59 @@ class IdentityReflectionTest(fixtures.TablesTest):
                 )
 
 
+class CompositeKeyReflectionTest(fixtures.TablesTest):
+    __backend__ = True
+
+    @classmethod
+    def define_tables(cls, metadata):
+        tb1 = Table(
+            "tb1",
+            metadata,
+            Column("id", Integer),
+            Column("attr", Integer),
+            Column("name", sql_types.VARCHAR(20)),
+            sa.PrimaryKeyConstraint("name", "id", "attr", name="pk_tb1"),
+            schema=None,
+            test_needs_fk=True,
+        )
+        Table(
+            "tb2",
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("pid", Integer),
+            Column("pattr", Integer),
+            Column("pname", sql_types.VARCHAR(20)),
+            sa.ForeignKeyConstraint(
+                ["pname", "pid", "pattr"],
+                [tb1.c.name, tb1.c.id, tb1.c.attr],
+                name="fk_tb1_name_id_attr",
+            ),
+            schema=None,
+            test_needs_fk=True,
+        )
+
+    @testing.requires.primary_key_constraint_reflection
+    @testing.provide_metadata
+    def test_pk_column_order(self):
+        # test for issue #5661
+        meta = self.metadata
+        insp = inspect(meta.bind)
+        primary_key = insp.get_pk_constraint(self.tables.tb1.name)
+        eq_(primary_key.get("constrained_columns"), ["name", "id", "attr"])
+
+    @testing.requires.foreign_key_constraint_reflection
+    @testing.provide_metadata
+    def test_fk_column_order(self):
+        # test for issue #5661
+        meta = self.metadata
+        insp = inspect(meta.bind)
+        foreign_keys = insp.get_foreign_keys(self.tables.tb2.name)
+        eq_(len(foreign_keys), 1)
+        fkey1 = foreign_keys[0]
+        eq_(fkey1.get("referred_columns"), ["name", "id", "attr"])
+        eq_(fkey1.get("constrained_columns"), ["pname", "pid", "pattr"])
+
+
 __all__ = (
     "ComponentReflectionTest",
     "QuotedNameArgumentTest",
@@ -1589,4 +1642,5 @@ __all__ = (
     "NormalizedNameTest",
     "ComputedReflectionTest",
     "IdentityReflectionTest",
+    "CompositeKeyReflectionTest",
 )