table.add_foreign_key(column, other_table, other_column)

New mechanism for adding foreign key constraints to an existing SQLite
table.

SQLite ALTER TABLE does not support this out-of-the-box, so I instead
had to write some careful code that uses PRAGMA writable_schema = 1 to
directly modify the sqlite_master table.

Refs #2

docs/python-api.rst

@@ -182,6 +182,7 @@ Note that the ``pk`` and ``column_order`` parameters here are optional if you ar
 
 An ``upsert_all()`` method is also available, which behaves like ``insert_all()`` but performs upserts instead.
 
+.. _python_api_add_column:
 
 Adding columns
 ==============
@@ -223,6 +224,31 @@ If you pass a Python type, it will be mapped to SQLite types as shown here::
     np.float32: "FLOAT"
     np.float64: "FLOAT"
 
+.. _python_api_add_foreign_key:
+
+Adding foreign key constraints
+==============================
+
+The SQLite ``ALTER TABLE`` statement doesn't have the ability to add foreign key references to an existing column.
+
+It's possible to add these references through very careful manipulation of SQLite's ``sqlite_master`` table, using ``PRAGMA writable_schema``.
+
+``sqlite-utils`` can do this for you, though there is a significant risk of data corruption if something goes wrong so it is advisable to create a fresh copy of your database file before attempting this.
+
+Here's an example of this mechanism in action:
+
+.. code-block:: python
+
+    db["authors"].insert_all([
+        {"id": 1, "name": "Sally"},
+        {"id": 2, "name": "Asheesh"}
+    ], pk="id")
+    db["books"].insert_all([
+        {"title": "Hedgehogs of the world", "author_id": 1},
+        {"title": "How to train your wolf", "author_id": 2},
+    ])
+    db["books"].add_foreign_key("author_id", "authors", "id")
+
 .. _python_api_hash:
 
 Setting an ID based on the hash of the row contents

sqlite_utils/db.py

@@ -59,6 +59,10 @@ if np:
     )
 
 
+class AlterError(Exception):
+    pass
+
+
 class Database:
     def __init__(self, filename_or_conn):
         if isinstance(filename_or_conn, str):
@@ -274,20 +278,42 @@ class Table:
     def drop(self):
         return self.db.conn.execute("DROP TABLE {}".format(self.name))
 
-    def add_foreign_key(self, column, column_type, other_table, other_column):
-        sql = """
-            ALTER TABLE {table} ADD COLUMN {column} {column_type}
-            REFERENCES {other_table}({other_column});
-        """.format(
-            table=self.name,
-            column=column,
-            column_type=column_type,
-            other_table=other_table,
-            other_column=other_column,
-        )
-        self.db.conn.execute(sql)
-        self.db.conn.commit()
-        return self
+    def add_foreign_key(self, column, other_table, other_column):
+        # Sanity check that the other column exists
+        if not [c for c in self.db[other_table].columns if c.name == other_column]:
+            raise AlterError("No such column: {}.{}".format(other_table, other_column))
+        # Check we do not already have an existing foreign key
+        if any(
+            fk
+            for fk in self.foreign_keys
+            if fk.column == column
+            and fk.other_table == other_table
+            and fk.other_column == other_column
+        ):
+            raise AlterError(
+                "Foreign key already exists for {} => {}.{}".format(
+                    column, other_table, other_column
+                )
+            )
+        with self.db.conn:
+            cursor = self.db.conn.cursor()
+            schema_version = cursor.execute("PRAGMA schema_version").fetchone()[0]
+            cursor.execute("PRAGMA writable_schema = 1")
+            old_sql = self.schema
+            extra_sql = ",\n   FOREIGN KEY({column}) REFERENCES {other_table}({other_column})\n".format(
+                column=column, other_table=other_table, other_column=other_column
+            )
+            # Stick that bit in at the very end just before the closing ')'
+            last_paren = old_sql.rindex(")")
+            new_sql = old_sql[:last_paren].strip() + extra_sql + old_sql[last_paren:]
+            cursor.execute(
+                "UPDATE sqlite_master SET sql = ? WHERE name = ?", (new_sql, self.name)
+            )
+            cursor.execute("PRAGMA schema_version = %d" % (schema_version + 1))
+            cursor.execute("PRAGMA writable_schema = 0")
+        # Have to VACUUM outside the transaction to ensure .foreign_keys property
+        # can see the newly created foreign key.
+        cursor.execute("VACUUM")
 
     def enable_fts(self, columns, fts_version="FTS5"):
         "Enables FTS on the specified columns"

tests/test_create.py

@@ -1,4 +1,4 @@
-from sqlite_utils.db import Index, Database
+from sqlite_utils.db import Index, Database, ForeignKey, AlterError
 import collections
 import datetime
 import json
@@ -150,6 +150,40 @@ def test_add_column(fresh_db, col_name, col_type, expected_schema):
     assert expected_schema == collapse_whitespace(fresh_db["dogs"].schema)
 
 
+def test_add_foreign_key(fresh_db):
+    fresh_db["authors"].insert_all(
+        [{"id": 1, "name": "Sally"}, {"id": 2, "name": "Asheesh"}], pk="id"
+    )
+    fresh_db["books"].insert_all(
+        [
+            {"title": "Hedgehogs of the world", "author_id": 1},
+            {"title": "How to train your wolf", "author_id": 2},
+        ]
+    )
+    assert [] == fresh_db["books"].foreign_keys
+    fresh_db["books"].add_foreign_key("author_id", "authors", "id")
+    assert [
+        ForeignKey(
+            table="books", column="author_id", other_table="authors", other_column="id"
+        )
+    ] == fresh_db["books"].foreign_keys
+
+
+def test_add_foreign_key_error_if_column_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")
+
+
+def test_add_foreign_key_error_if_already_exists(fresh_db):
+    fresh_db["books"].insert({"title": "Hedgehogs of the world", "author_id": 1})
+    fresh_db["authors"].insert({"id": 1, "name": "Sally"}, pk="id")
+    fresh_db["books"].add_foreign_key("author_id", "authors", "id")
+    with pytest.raises(AlterError) as ex:
+        fresh_db["books"].add_foreign_key("author_id", "authors", "id")
+    assert "Foreign key already exists for author_id => authors.id" == ex.value.args[0]
+
+
 @pytest.mark.parametrize(
     "columns,index_name,expected_index",
     (