Compound primary keys now resolve in PRIMARY KEY declaration order

PRAGMA table_info sets is_pk to the 1-based position of each column
within the PRIMARY KEY, which can differ from table column order.
table.pks previously returned table column order, so an implicit
compound FOREIGN KEY ... REFERENCES other was introspected with its
referenced columns inverted, and transform() baked that inverted order
into the rewritten schema - failing with IntegrityError on valid data,
or silently reversing the constraint with foreign keys off.

table.pks, compound foreign key guessing (create, add_foreign_key) and
transform() now all use declaration order, and transform() no longer
reorders a compound PRIMARY KEY (b, a) into table column order.

Refs https://github.com/simonw/sqlite-utils/issues/769#issuecomment-4900034150

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

docs/changelog.rst

@@ -19,6 +19,7 @@ Unreleased
 - Fixed a bug where a failed write statement executed with ``db.execute()`` left the driver's implicit transaction open. Every subsequent write then joined that phantom transaction, which nothing committed, so their work was silently rolled back when the connection was closed. The implicit transaction opened by a failed statement is now rolled back before the exception is raised. A failed write inside a transaction opened with ``db.begin()`` or ``db.atomic()`` leaves that transaction open and untouched, as before.
 - Fixed a bug where transaction-control statements prefixed with an empty statement - ``db.query("; COMMIT")`` - or a UTF-8 byte order mark slipped past the check that rejects them, committing the caller's open transaction before raising a confusing ``OperationalError``. The keyword scanner used by ``db.query()`` and ``db.execute()`` now skips leading ``;`` and byte order marks, matching what the ``sqlite3`` driver tolerates before the first token, so these statements are rejected with a ``ValueError`` without being executed. The same fix means ``db.execute("; BEGIN")`` no longer auto-commits the transaction it just opened.
 - Documented a limitation of ``db.query()``: a ``PRAGMA`` statement that returns no rows raises a ``ValueError`` but still takes effect, because PRAGMA statements run outside the savepoint guard used to roll back other rejected statements. Use ``db.execute()`` for row-less PRAGMA statements.
+- Fixed a bug where compound primary key columns were returned in table column order instead of ``PRIMARY KEY`` declaration order. For a table declared as ``CREATE TABLE other (b TEXT, a TEXT, PRIMARY KEY (a, b))`` an implicit ``FOREIGN KEY (x, y) REFERENCES other`` was introspected as referencing ``(b, a)`` when SQLite resolves it as ``(a, b)`` - running ``transform()`` on such a table then rewrote the schema with the inverted column order, silently reversing the meaning of the constraint and causing foreign key errors on valid data. ``table.pks``, compound foreign key guessing and ``transform()`` now all use the primary key declaration order, and ``transform()`` no longer reorders a compound ``PRIMARY KEY (b, a)`` into table column order.
 
 .. _v4_0rc3:
 

sqlite_utils/db.py

@@ -2012,10 +2012,9 @@ class Queryable:
         :param offset: Integer for SQL offset
         """
         column_names = [column.name for column in self.columns]
-        pks = [column.name for column in self.columns if column.is_pk]
-        if not pks:
+        pks = self.pks
+        if self.use_rowid:
             column_names.insert(0, "rowid")
-            pks = ["rowid"]
         select = ",".join(quote_identifier(column_name) for column_name in column_names)
         for row in self.rows_where(
             select=select,
@@ -2148,8 +2147,18 @@ class Table(Queryable):
 
     @property
     def pks(self) -> List[str]:
-        "Primary key columns for this table."
-        names = [column.name for column in self.columns if column.is_pk]
+        """
+        Primary key columns for this table, in PRIMARY KEY declaration order -
+        ``PRAGMA table_info`` sets ``is_pk`` to the 1-based position of each
+        column within the primary key, which can differ from the order of the
+        columns in the table. SQLite uses the declaration order to resolve
+        implicit foreign key references, so this order matters.
+        """
+        pk_columns = sorted(
+            (column for column in self.columns if column.is_pk),
+            key=lambda column: column.is_pk,
+        )
+        names = [column.name for column in pk_columns]
         if not names:
             names = ["rowid"]
         return names
@@ -2673,7 +2682,8 @@ class Table(Queryable):
 
         if pk is DEFAULT:
             pks_renamed = tuple(
-                rename.get(p.name) or p.name for p in self.columns if p.is_pk
+                rename.get(pk_name) or pk_name
+                for pk_name in (self.pks if not self.use_rowid else [])
             )
             if len(pks_renamed) == 1:
                 pk = pks_renamed[0]
@@ -3017,7 +3027,10 @@ class Table(Queryable):
                     raise AlterError("table '{}' has no column {}".format(fk, fk_col))
             else:
                 # automatically set fk_col to first primary_key of fk table
-                pks = [c for c in self.db[fk].columns if c.is_pk]
+                pks = sorted(
+                    (c for c in self.db[fk].columns if c.is_pk),
+                    key=lambda c: c.is_pk,
+                )
                 if pks:
                     fk_col = pks[0].name
                     fk_col_type = pks[0].type
@@ -3770,9 +3783,7 @@ class Table(Queryable):
         # First we execute the function
         pk_to_values = {}
         new_column_types: Dict[str, Set[type]] = {}
-        pks = [column.name for column in self.columns if column.is_pk]
-        if not pks:
-            pks = ["rowid"]
+        pks = self.pks
 
         with progressbar(
             length=self.count, silent=not show_progress, label="1: Evaluating"

tests/test_foreign_keys.py

@@ -524,3 +524,57 @@ def test_add_compound_foreign_key_on_delete(courses_db):
     assert fk.is_compound is True
     assert fk.on_delete == "SET NULL"
     assert "ON DELETE SET NULL" in courses_db["courses"].schema
+
+
+def test_implicit_compound_foreign_key_resolves_pk_declaration_order(fresh_db):
+    # The other table's PRIMARY KEY declares its columns in a different
+    # order to the table's column order. SQLite resolves the implicit
+    # "REFERENCES other" using PRIMARY KEY declaration order, so the
+    # introspected other_columns must too
+    fresh_db.execute("create table other (b text, a text, primary key (a, b))")
+    fresh_db.execute(
+        "create table child (x text, y text, foreign key (x, y) references other)"
+    )
+    fk = fresh_db["child"].foreign_keys[0]
+    assert fk.other_columns == ("a", "b")
+
+
+def test_transform_implicit_compound_foreign_key_stays_valid(fresh_db):
+    # transform() rewrites the implicit FK with explicit columns - they
+    # must be in PRIMARY KEY declaration order or valid data fails the
+    # foreign key check with an IntegrityError
+    fresh_db.execute("create table other (b text, a text, primary key (a, b))")
+    fresh_db.execute(
+        "create table child (x text, y text, foreign key (x, y) references other)"
+    )
+    fresh_db.execute("PRAGMA foreign_keys = ON")
+    fresh_db["other"].insert({"a": "A", "b": "B"})
+    fresh_db["child"].insert({"x": "A", "y": "B"})
+    fresh_db["child"].transform(types={"x": str})
+    assert fresh_db["child"].foreign_keys[0].other_columns == ("a", "b")
+    # The constraint still points the right way around
+    fresh_db["child"].insert({"x": "A", "y": "B"})
+    with pytest.raises(sqlite3.IntegrityError):
+        fresh_db["child"].insert({"x": "B", "y": "A"})
+
+
+def test_create_compound_foreign_key_guesses_pk_declaration_order(fresh_db):
+    fresh_db.execute("create table other (b text, a text, primary key (a, b))")
+    fresh_db["other"].insert({"a": "A", "b": "B"})
+    fresh_db["child"].create(
+        {"id": int, "x": str, "y": str},
+        pk="id",
+        foreign_keys=[(("x", "y"), "other")],
+    )
+    assert fresh_db["child"].foreign_keys[0].other_columns == ("a", "b")
+    fresh_db.execute("PRAGMA foreign_keys = ON")
+    fresh_db["child"].insert({"id": 1, "x": "A", "y": "B"})
+    with pytest.raises(sqlite3.IntegrityError):
+        fresh_db["child"].insert({"id": 2, "x": "B", "y": "A"})
+
+
+def test_add_compound_foreign_key_guesses_pk_declaration_order(fresh_db):
+    fresh_db.execute("create table other (b text, a text, primary key (a, b))")
+    fresh_db["child"].insert({"id": 1, "x": "A", "y": "B"}, pk="id")
+    fresh_db["child"].add_foreign_key(("x", "y"), "other")
+    assert fresh_db["child"].foreign_keys[0].other_columns == ("a", "b")

tests/test_introspect.py

@@ -321,3 +321,18 @@ def test_table_default_values(fresh_db, value):
     )
     default_values = fresh_db["default_values"].default_values
     assert default_values == {"value": value}
+
+
+def test_pks_use_primary_key_declaration_order(fresh_db):
+    # PRIMARY KEY (a, b) declared against columns stored in order (b, a) -
+    # pks must follow the declaration order, which is what SQLite uses to
+    # resolve implicit foreign key references and compound pk lookups
+    fresh_db.execute("create table t (b text, a text, primary key (a, b))")
+    assert fresh_db["t"].pks == ["a", "b"]
+
+
+def test_transform_preserves_compound_pk_declaration_order(fresh_db):
+    fresh_db.execute("create table t (a text, b text, c text, primary key (b, a))")
+    fresh_db["t"].transform(drop={"c"})
+    assert fresh_db["t"].pks == ["b", "a"]
+    assert 'PRIMARY KEY ("b", "a")' in fresh_db["t"].schema