diff --git a/docs/python-api.rst b/docs/python-api.rst index 9ae6646..1ffc94d 100644 --- a/docs/python-api.rst +++ b/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 diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index c1bbd74..c30de6f 100644 --- a/sqlite_utils/db.py +++ b/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" diff --git a/tests/test_create.py b/tests/test_create.py index 29618dc..4ccfcc0 100644 --- a/tests/test_create.py +++ b/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", (