db.add_foreign_keys() method

Closes #31

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

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"

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