diff --git a/docs/python-api.rst b/docs/python-api.rst index d1f466f..5e55256 100644 --- a/docs/python-api.rst +++ b/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 diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 67623fe..12914d1 100644 --- a/sqlite_utils/db.py +++ b/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): - if len(tuple_or_list) == 4: - tuple_or_list = cast(Tuple[str, str, str], tuple_or_list[1:]) + column_or_columns = tuple_or_list[0] + other_table = 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) - fks.append( - ForeignKey( - name, tuple_or_list[0], tuple_or_list[1], tuple_or_list[2] + # Guess the compound primary key of the other table + other_columns = self.table(other_table).pks + if len(columns) != len(other_columns): + 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], - tuple_or_list[1], - table.guess_foreign_column(tuple_or_list[1]), + column_or_columns, + other_table, + 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)): - continue - if fk.other_column != "rowid" and not any( - c for c in self[fk.other_table].columns if c.name == fk.other_column - ): - raise AlterError( - "No such column: {}.{}".format(fk.other_table, fk.other_column) - ) + for other_column in fk.other_columns: + if fk.other_table == name and columns.get(other_column): + continue + 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, 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} diff --git a/tests/test_foreign_keys.py b/tests/test_foreign_keys.py index 4b371a9..457f307 100644 --- a/tests/test_foreign_keys.py +++ b/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"]) + ], + )