diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index ec9ec95..1b58f94 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -466,6 +466,9 @@ class Table: return pks[0].name def add_foreign_key(self, column, other_table=None, other_column=None): + # Ensure column exists + if column not in self.columns_dict: + raise AlterError("No such column: {}".format(column)) # If other_table is not specified, attempt to guess it from the column if other_table is None: other_table = self.guess_foreign_table(column) diff --git a/tests/test_create.py b/tests/test_create.py index 27c1a63..a275a7f 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -280,6 +280,14 @@ def test_add_foreign_key(fresh_db): def test_add_foreign_key_error_if_column_does_not_exist(fresh_db): + fresh_db["books"].insert( + {"id": 1, "title": "Hedgehogs of the world", "author_id": 1} + ) + with pytest.raises(AlterError): + fresh_db["books"].add_foreign_key("author2_id", "books", "id") + + +def test_add_foreign_key_error_if_other_table_does_not_exist(fresh_db): fresh_db["books"].insert({"title": "Hedgehogs of the world", "author_id": 1}) with pytest.raises(AlterError): fresh_db["books"].add_foreign_key("author_id", "authors", "id")