Sanity check add_foreign_key() column exists, closes #29

This commit is contained in:
Simon Willison 2019-06-20 16:58:09 -07:00 committed by Simon Willison
commit 976263effc
2 changed files with 11 additions and 0 deletions

View file

@ -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)

View file

@ -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")