From 997d8758fc8992be93d5984556e5b9e3bad4205f Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Fri, 28 Jun 2019 23:27:38 -0700 Subject: [PATCH] db.add_foreign_keys() method Closes #31 --- docs/python-api.rst | 20 +++++++++++ sqlite_utils/db.py | 83 ++++++++++++++++++++++++++++++++++---------- tests/test_create.py | 28 +++++++++++++++ 3 files changed, 112 insertions(+), 19 deletions(-) diff --git a/docs/python-api.rst b/docs/python-api.rst index 1243a17..8928c03 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -399,6 +399,26 @@ The ``table.add_foreign_key(column, other_table, other_column)`` method takes th - If the column is of format ``author_id``, look for tables called ``author`` or ``authors`` - If the column does not end in ``_id``, try looking for a table with the exact name of the column or that name with an added ``s`` +.. _python_api_add_foreign_keys: + +Adding multiple foreign key constraints at once +----------------------------------------------- + +The final step in adding a new foreign key to a SQLite database is to run ``VACUUM``, to ensure the new foreign key is available in future introspection queries. + +``VACUUM`` against a large (multi-GB) database can take several minutes or longer. If you are adding multiple foreign keys using ``table.add_foreign_key(...)`` these can quickly add up. + +Instead, you can use ``db.add_foreign_keys(...)`` to add multiple foreign keys within a single transaction. This method takes a list of four-tuples, each one specifying a ``table``, ``column``, ``other_table`` and ``other_column``. + +Here's an example adding two foreign keys at once: + +.. code-block:: python + + db.add_foreign_keys([ + ("dogs", "breed_id", "breeds", "id"), + ("dogs", "home_town_id", "towns", "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 414ec97..582017f 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -252,6 +252,69 @@ class Database: ) ) + def add_foreign_keys(self, foreign_keys): + # foreign_keys is a list of explicit 4-tuples + assert all( + len(fk) == 4 and isinstance(fk, (list, tuple)) for fk in foreign_keys + ), "foreign_keys must be a list of 4-tuples, (table, column, other_table, other_column)" + + foreign_keys_to_create = [] + + # Verify that all tables and columns exist + for table, column, other_table, other_column in foreign_keys: + if not self[table].exists: + raise AlterError("No such table: {}".format(table)) + if column not in self[table].columns_dict: + raise AlterError("No such column: {} in {}".format(column, table)) + if not self[other_table].exists: + raise AlterError("No such other_table: {}".format(other_table)) + if ( + other_column != "rowid" + and other_column not in self[other_table].columns_dict + ): + raise AlterError( + "No such other_column: {} in {}".format(other_column, other_table) + ) + # We will silently skip foreign keys that exist already + if not any( + fk + for fk in self[table].foreign_keys + if fk.column == column + and fk.other_table == other_table + and fk.other_column == other_column + ): + foreign_keys_to_create.append( + (table, column, other_table, other_column) + ) + + # Construct SQL for use with "UPDATE sqlite_master SET sql = ? WHERE name = ?" + table_sql = {} + for table, column, other_table, other_column in foreign_keys_to_create: + old_sql = table_sql.get(table, self[table].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:] + table_sql[table] = new_sql + + # And execute it all within a single transaction + with self.conn: + cursor = self.conn.cursor() + schema_version = cursor.execute("PRAGMA schema_version").fetchone()[0] + cursor.execute("PRAGMA writable_schema = 1") + for table_name, new_sql in table_sql.items(): + cursor.execute( + "UPDATE sqlite_master SET sql = ? WHERE name = ?", + (new_sql, table_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. + self.vacuum() + def vacuum(self): self.conn.execute("VACUUM;") @@ -495,25 +558,7 @@ class Table: 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") + self.db.add_foreign_keys([(self.name, column, other_table, other_column)]) 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 04b1fbf..88d7676 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -303,6 +303,34 @@ def test_add_foreign_key_error_if_already_exists(fresh_db): assert "Foreign key already exists for author_id => authors.id" == ex.value.args[0] +def test_add_foreign_keys(fresh_db): + fresh_db["authors"].insert_all( + [{"id": 1, "name": "Sally"}, {"id": 2, "name": "Asheesh"}], pk="id" + ) + fresh_db["categories"].insert_all([{"id": 1, "name": "Wildlife"}], pk="id") + fresh_db["books"].insert_all( + [{"title": "Hedgehogs of the world", "author_id": 1, "category_id": 1}] + ) + assert [] == fresh_db["books"].foreign_keys + fresh_db.add_foreign_keys( + [ + ("books", "author_id", "authors", "id"), + ("books", "category_id", "categories", "id"), + ] + ) + assert [ + ForeignKey( + table="books", column="author_id", other_table="authors", other_column="id" + ), + ForeignKey( + table="books", + column="category_id", + other_table="categories", + other_column="id", + ), + ] == sorted(fresh_db["books"].foreign_keys) + + def test_add_column_foreign_key(fresh_db): fresh_db.create_table("dogs", {"name": str}) fresh_db.create_table("breeds", {"name": str})