Sanity check add_foreign_key() column exists, closes #29

Author: simonw

sqlite_utils/db.py

@@ -466,6 +466,9 @@ def guess_foreign_column(self, other_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)

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