transform() now round-trips compound foreign keys, refs #594

Compound foreign keys are passed through table.transform() intact
instead of being degraded to per-column single foreign keys:

- rename= applies to each member column of a compound key
- drop= of any member column drops the whole constraint, matching
  the existing single-column behavior
- drop_foreign_keys= accepts a bare column name (drops any foreign
  key that column participates in) or a tuple of columns (drops the
  compound key with exactly those columns)

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

docs/python-api.rst

@@ -1793,6 +1793,16 @@ This example drops two foreign keys - the one from ``places.country`` to ``count
         drop_foreign_keys=("country", "continent")
     )
 
+A bare column name drops any foreign key that column participates in, including compound foreign keys. To target a compound foreign key precisely, pass a tuple of its columns:
+
+.. code-block:: python
+
+    db.table("courses").transform(
+        drop_foreign_keys=[("campus_name", "dept_code")]
+    )
+
+Renaming a column with ``rename=`` updates any foreign keys that use it, and dropping a column with ``drop=`` also drops any foreign keys it participates in - for a compound foreign key this removes the whole constraint.
+
 .. _python_api_transform_sql:
 
 Custom transformations with .transform_sql()

sqlite_utils/db.py

@@ -2275,7 +2275,9 @@ class Table(Queryable):
         :param pk: New primary key for the table
         :param not_null: Columns to set as ``NOT NULL``
         :param defaults: Default values for columns
-        :param drop_foreign_keys: Names of columns that should have their foreign key constraints removed
+        :param drop_foreign_keys: Foreign key constraints to remove - a column name
+          drops any foreign key that column participates in, a tuple of column names
+          drops the compound foreign key with exactly those columns
         :param add_foreign_keys: List of foreign keys to add to the table
         :param foreign_keys: List of foreign keys to set for the table, replacing any existing foreign keys
         :param column_order: List of strings specifying a full or partial column order
@@ -2360,7 +2362,9 @@ class Table(Queryable):
         :param pk: New primary key for the table
         :param not_null: Columns to set as ``NOT NULL``
         :param defaults: Default values for columns
-        :param drop_foreign_keys: Names of columns that should have their foreign key constraints removed
+        :param drop_foreign_keys: Foreign key constraints to remove - a column name
+          drops any foreign key that column participates in, a tuple of column names
+          drops the compound foreign key with exactly those columns
         :param add_foreign_keys: List of foreign keys to add to the table
         :param foreign_keys: List of foreign keys to set for the table, replacing any existing foreign keys
         :param column_order: List of strings specifying a full or partial column order
@@ -2387,32 +2391,45 @@ class Table(Queryable):
             create_table_foreign_keys.extend(foreign_keys)
         else:
             # Construct foreign_keys from current, plus add_foreign_keys, minus drop_foreign_keys
+
+            def fk_should_be_dropped(fk: ForeignKey) -> bool:
+                if drop_foreign_keys is not None:
+                    for spec in drop_foreign_keys:
+                        if isinstance(spec, str):
+                            # A column name matches any foreign key it participates in
+                            if spec in fk.columns:
+                                return True
+                        elif list(spec) == fk.columns:
+                            # A tuple/list must match a compound key's columns exactly
+                            return True
+                # Dropping any of a foreign key's columns drops the whole key
+                return any(column in drop for column in fk.columns)
+
+            def fk_with_renamed_columns(fk: ForeignKey) -> ForeignKey:
+                columns = [rename.get(column) or column for column in fk.columns]
+                if fk.is_compound:
+                    return ForeignKey(
+                        self.name,
+                        None,
+                        fk.other_table,
+                        None,
+                        columns=columns,
+                        other_columns=fk.other_columns,
+                        is_compound=True,
+                    )
+                return ForeignKey(
+                    self.name, columns[0], fk.other_table, fk.other_columns[0]
+                )
+
             create_table_foreign_keys = []
+            # Copy over old foreign keys, unless we are dropping them
             for fk in self.foreign_keys:
-                # Expand compound foreign keys into per-column references;
-                # for single-column keys this iterates exactly once
-                for column, other_column in zip(fk.columns, fk.other_columns):
-                    # Copy over old foreign keys, unless we are dropping them
-                    if (drop_foreign_keys is None) or (column not in drop_foreign_keys):
-                        create_table_foreign_keys.append(
-                            ForeignKey(
-                                fk.table,
-                                rename.get(column) or column,
-                                fk.other_table,
-                                other_column,
-                            )
-                        )
+                if not fk_should_be_dropped(fk):
+                    create_table_foreign_keys.append(fk_with_renamed_columns(fk))
             # Add new foreign keys
             if add_foreign_keys is not None:
                 for fk in self.db.resolve_foreign_keys(self.name, add_foreign_keys):
-                    create_table_foreign_keys.append(
-                        ForeignKey(
-                            self.name,
-                            rename.get(fk.column) or fk.column,
-                            fk.other_table,
-                            fk.other_column,
-                        )
-                    )
+                    create_table_foreign_keys.append(fk_with_renamed_columns(fk))
 
         new_table_name = "{}_new_{}".format(
             self.name, tmp_suffix or os.urandom(6).hex()

tests/test_foreign_keys.py

@@ -202,3 +202,51 @@ def test_create_table_compound_foreign_key_missing_other_column(departments_db):
                 (["campus_name", "dept_code"], "departments", ["campus_name", "nope"])
             ],
         )
+
+
+def test_transform_preserves_compound_foreign_key(compound_db):
+    compound_db["courses"].transform(rename={"course_name": "title"})
+    fks = compound_db["courses"].foreign_keys
+    assert len(fks) == 1
+    fk = fks[0]
+    assert fk.is_compound is True
+    assert fk.columns == ["campus_name", "dept_code"]
+    assert fk.other_table == "departments"
+    assert fk.other_columns == ["campus_name", "dept_code"]
+
+
+def test_transform_rename_member_column_updates_compound_foreign_key(compound_db):
+    compound_db["courses"].transform(rename={"campus_name": "campus"})
+    fks = compound_db["courses"].foreign_keys
+    assert len(fks) == 1
+    fk = fks[0]
+    assert fk.is_compound is True
+    assert fk.columns == ["campus", "dept_code"]
+    # Referenced columns in the other table are unchanged
+    assert fk.other_columns == ["campus_name", "dept_code"]
+
+
+def test_transform_drop_member_column_drops_compound_foreign_key(compound_db):
+    # Matches single-column behavior: dropping the column silently
+    # drops the foreign key that used it
+    compound_db["courses"].transform(drop={"dept_code"})
+    assert compound_db["courses"].foreign_keys == []
+    assert "FOREIGN KEY" not in compound_db["courses"].schema
+
+
+@pytest.mark.parametrize(
+    "drop_foreign_keys",
+    (
+        # A bare column name matches any foreign key it participates in:
+        ["campus_name"],
+        # A tuple must match the full compound key:
+        [("campus_name", "dept_code")],
+    ),
+)
+def test_transform_drop_compound_foreign_key(compound_db, drop_foreign_keys):
+    compound_db["courses"].transform(drop_foreign_keys=drop_foreign_keys)
+    assert compound_db["courses"].foreign_keys == []
+    # The columns themselves survive
+    assert {"campus_name", "dept_code"} <= set(
+        compound_db["courses"].columns_dict.keys()
+    )