Create tables with compound foreign keys, refs #594

create_table() and friends now accept compound foreign keys, specified
as lists of column names in the existing tuple forms:

    foreign_keys=[
        (["campus_name", "dept_code"], "departments"),
        (["campus_name", "dept_code"], "departments", ["campus_name", "dept_code"]),
    ]

The two-item form guesses the compound primary key of the other table.
Compound keys are rendered as table-level FOREIGN KEY constraints,
while single-column keys keep their inline REFERENCES clauses.

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

docs/python-api.rst

@@ -817,6 +817,42 @@ You can leave off the third item in the tuple to have the referenced column auto
         ("author_id", "authors")
     ])
 
+.. _python_api_compound_foreign_keys:
+
+Compound foreign keys
+~~~~~~~~~~~~~~~~~~~~~
+
+To create a compound (multi-column) foreign key, use lists of column names in place of the single column names:
+
+.. code-block:: python
+
+    db.table("courses").create({
+        "course_code": str,
+        "campus_name": str,
+        "dept_code": str,
+    }, pk="course_code", foreign_keys=[
+        (["campus_name", "dept_code"], "departments", ["campus_name", "dept_code"])
+    ])
+
+This creates a table-level constraint:
+
+.. code-block:: sql
+
+    CREATE TABLE "courses" (
+       "course_code" TEXT PRIMARY KEY,
+       "campus_name" TEXT,
+       "dept_code" TEXT,
+       FOREIGN KEY ("campus_name", "dept_code") REFERENCES "departments"("campus_name", "dept_code")
+    )
+
+As with single columns, you can leave off the list of other columns to reference the compound primary key of the other table:
+
+.. code-block:: python
+
+    foreign_keys=[
+        (["campus_name", "dept_code"], "departments")
+    ]
+
 .. _python_api_table_configuration:
 
 Table configuration options

sqlite_utils/db.py

@@ -219,6 +219,10 @@ ForeignKeyIndicator = Union[
     Tuple[str, str],
     Tuple[str, str, str],
     Tuple[str, str, str, str],
+    # Compound foreign keys use lists of columns:
+    Tuple[List[str], str],
+    Tuple[List[str], str, List[str]],
+    Tuple[str, List[str], str, List[str]],
 ]
 
 ForeignKeysType = Union[Iterable[ForeignKeyIndicator], List[ForeignKeyIndicator]]
@@ -1142,9 +1146,12 @@ class Database:
 
         :param name: Name of table that foreign keys are being defined for
         :param foreign_keys: List of foreign keys, each of which can be a
-            string, a ForeignKey() named tuple, a tuple of (column, other_table),
+            string, a ForeignKey() object, a tuple of (column, other_table),
             or a tuple of (column, other_table, other_column), or a tuple of
-            (table, column, other_table, other_column)
+            (table, column, other_table, other_column). For compound foreign
+            keys the column elements can be lists of column names, e.g.
+            (["campus_name", "dept_code"], "departments") or
+            (["campus_name", "dept_code"], "departments", ["campus_name", "dept_code"])
         """
         table = self.table(name)
         if all(isinstance(fk, ForeignKey) for fk in foreign_keys):
@@ -1161,7 +1168,7 @@ class Database:
         if not all(isinstance(fk, (tuple, list)) for fk in foreign_keys):
             raise ValueError("foreign_keys= should be a list of tuples")
         fks = []
-        for tuple_or_list in cast(Iterable[Sequence[str]], foreign_keys):
+        for tuple_or_list in cast(Iterable[Sequence[Any]], foreign_keys):
             if len(tuple_or_list) == 4:
                 if tuple_or_list[0] != name:
                     raise ValueError(
@@ -1169,28 +1176,61 @@ class Database:
                             tuple_or_list, name
                         )
                     )
-            if len(tuple_or_list) not in (2, 3, 4):
+                tuple_or_list = tuple_or_list[1:]
+            if len(tuple_or_list) not in (2, 3):
                 raise ValueError(
                     "foreign_keys= should be a list of tuple pairs or triples"
                 )
-            if len(tuple_or_list) in (3, 4):
+            column_or_columns = tuple_or_list[0]
-                if len(tuple_or_list) == 4:
+            other_table = tuple_or_list[1]
-                    tuple_or_list = cast(Tuple[str, str, str], tuple_or_list[1:])
+            if isinstance(column_or_columns, (list, tuple)):
+                # Compound foreign key
+                columns = list(column_or_columns)
+                if len(tuple_or_list) == 3:
+                    other_columns = tuple_or_list[2]
+                    if not isinstance(other_columns, (list, tuple)):
+                        raise ValueError(
+                            "Compound foreign key {} should reference a list "
+                            "of other columns".format(tuple(tuple_or_list))
+                        )
+                    other_columns = list(other_columns)
                 else:
-                    tuple_or_list = cast(Tuple[str, str, str], tuple_or_list)
+                    # Guess the compound primary key of the other table
-                fks.append(
+                    other_columns = self.table(other_table).pks
-                    ForeignKey(
+                if len(columns) != len(other_columns):
-                        name, tuple_or_list[0], tuple_or_list[1], tuple_or_list[2]
+                    raise ValueError(
+                        "Compound foreign key {} should have the same number "
+                        "of columns on both sides".format(tuple(tuple_or_list))
                     )
+                if len(columns) == 1:
+                    # Single-column key passed as a one-item list
+                    fks.append(
+                        ForeignKey(name, columns[0], other_table, other_columns[0])
+                    )
+                else:
+                    fks.append(
+                        ForeignKey(
+                            name,
+                            None,
+                            other_table,
+                            None,
+                            columns=columns,
+                            other_columns=other_columns,
+                            is_compound=True,
+                        )
+                    )
+            elif len(tuple_or_list) == 3:
+                fks.append(
+                    ForeignKey(name, column_or_columns, other_table, tuple_or_list[2])
                 )
             else:
                 # Guess the primary key
                 fks.append(
                     ForeignKey(
                         name,
-                        tuple_or_list[0],
+                        column_or_columns,
-                        tuple_or_list[1],
+                        other_table,
-                        table.guess_foreign_column(tuple_or_list[1]),
+                        table.guess_foreign_column(other_table),
                     )
                 )
         return fks
@@ -1229,7 +1269,11 @@ class Database:
         if hash_id_columns and (hash_id is None):
             hash_id = "id"
         foreign_keys = self.resolve_foreign_keys(name, foreign_keys or [])
-        foreign_keys_by_column = {fk.column: fk for fk in foreign_keys}
+        # Compound foreign keys are rendered as table-level constraints;
+        # single-column ones as inline REFERENCES on their column
+        foreign_keys_by_column = {
+            fk.column: fk for fk in foreign_keys if not fk.is_compound
+        }
         # any extracts will be treated as integer columns with a foreign key
         extracts = resolve_extracts(extracts)
         for extract_column, extract_table in extracts.items():
@@ -1271,14 +1315,15 @@ class Database:
             pk = hash_id
         # Soundness check foreign_keys point to existing tables
         for fk in foreign_keys:
-            if fk.other_table == name and columns.get(cast(str, fk.other_column)):
+            for other_column in fk.other_columns:
-                continue
+                if fk.other_table == name and columns.get(other_column):
-            if fk.other_column != "rowid" and not any(
+                    continue
-                c for c in self[fk.other_table].columns if c.name == fk.other_column
+                if other_column != "rowid" and not any(
-            ):
+                    c for c in self[fk.other_table].columns if c.name == other_column
-                raise AlterError(
+                ):
-                    "No such column: {}.{}".format(fk.other_table, fk.other_column)
+                    raise AlterError(
-                )
+                        "No such column: {}.{}".format(fk.other_table, other_column)
+                    )
 
         column_defs = []
         # ensure pk is a tuple
@@ -1329,6 +1374,25 @@ class Database:
             extra_pk = ",\n   PRIMARY KEY ({pks})".format(
                 pks=", ".join([quote_identifier(p) for p in pk])
             )
+        # Compound foreign keys become table-level FOREIGN KEY constraints
+        column_names = [c[0] for c in column_items]
+        for fk in foreign_keys:
+            if not fk.is_compound:
+                continue
+            missing = [c for c in fk.columns if c not in column_names]
+            if missing:
+                raise AlterError(
+                    "No such column: {}".format(", ".join(sorted(missing)))
+                )
+            column_defs.append(
+                "   FOREIGN KEY ({columns}) REFERENCES {other_table}({other_columns})".format(
+                    columns=", ".join(quote_identifier(c) for c in fk.columns),
+                    other_table=quote_identifier(fk.other_table),
+                    other_columns=", ".join(
+                        quote_identifier(c) for c in fk.other_columns
+                    ),
+                )
+            )
         columns_sql = ",\n".join(column_defs)
         sql = """CREATE TABLE {if_not_exists}{table} (
 {columns_sql}{extra_pk}

tests/test_foreign_keys.py

@@ -1,7 +1,8 @@
-"""Tests for reading compound (multi-column) foreign keys - issue #594."""
+"""Tests for compound (multi-column) foreign keys - issue #594."""
 
 import pytest
 from sqlite_utils import Database
+from sqlite_utils.db import AlterError, ForeignKey
 
 COMPOUND_SCHEMA = """
 CREATE TABLE departments (
@@ -108,3 +109,96 @@ def test_mixed_compound_and_single_foreign_keys_are_sortable():
     fks_sorted = sorted(fks)
     assert fks_sorted[0].other_table == "accreditations"
     assert fks_sorted[1].other_table == "departments"
+
+
+@pytest.fixture
+def departments_db():
+    db = Database(memory=True)
+    db.create_table(
+        "departments",
+        {"campus_name": str, "dept_code": str, "dept_name": str},
+        pk=("campus_name", "dept_code"),
+    )
+    return db
+
+
+EXPECTED_COURSES_SCHEMA = (
+    'CREATE TABLE "courses" (\n'
+    '   "course_code" TEXT PRIMARY KEY,\n'
+    '   "campus_name" TEXT,\n'
+    '   "dept_code" TEXT,\n'
+    '   FOREIGN KEY ("campus_name", "dept_code") '
+    'REFERENCES "departments"("campus_name", "dept_code")\n'
+    ")"
+)
+
+
+@pytest.mark.parametrize(
+    "foreign_keys",
+    (
+        [
+            ForeignKey(
+                table="courses",
+                column=None,
+                other_table="departments",
+                other_column=None,
+                columns=["campus_name", "dept_code"],
+                other_columns=["campus_name", "dept_code"],
+                is_compound=True,
+            )
+        ],
+        [(["campus_name", "dept_code"], "departments", ["campus_name", "dept_code"])],
+        # Two-item form guesses the other table's primary key:
+        [(["campus_name", "dept_code"], "departments")],
+    ),
+)
+def test_create_table_with_compound_foreign_key(departments_db, foreign_keys):
+    departments_db.create_table(
+        "courses",
+        {"course_code": str, "campus_name": str, "dept_code": str},
+        pk="course_code",
+        foreign_keys=foreign_keys,
+    )
+    assert departments_db["courses"].schema == EXPECTED_COURSES_SCHEMA
+    fks = departments_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_create_table_compound_foreign_key_enforced(departments_db):
+    departments_db.execute("PRAGMA foreign_keys = ON")
+    departments_db.create_table(
+        "courses",
+        {"course_code": str, "campus_name": str, "dept_code": str},
+        pk="course_code",
+        foreign_keys=[(["campus_name", "dept_code"], "departments")],
+    )
+    departments_db["departments"].insert(
+        {"campus_name": "Berkeley", "dept_code": "CS", "dept_name": "Computer Science"}
+    )
+    departments_db["courses"].insert(
+        {"course_code": "CS101", "campus_name": "Berkeley", "dept_code": "CS"}
+    )
+    import sqlite3
+
+    with pytest.raises(sqlite3.IntegrityError):
+        departments_db.execute(
+            "insert into courses (course_code, campus_name, dept_code) "
+            "values ('X1', 'Nowhere', 'NOPE')"
+        )
+
+
+def test_create_table_compound_foreign_key_missing_other_column(departments_db):
+    with pytest.raises(AlterError):
+        departments_db.create_table(
+            "courses",
+            {"course_code": str, "campus_name": str, "dept_code": str},
+            pk="course_code",
+            foreign_keys=[
+                (["campus_name", "dept_code"], "departments", ["campus_name", "nope"])
+            ],
+        )