From 3cc1944e53b75749644f558cbe1717397cae72ea Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 20 Sep 2020 13:14:25 -0700
Subject: [PATCH] sqlite-utils add-foreign-keys command, closes #157

---
 docs/cli.rst        | 13 +++++++++++++
 sqlite_utils/cli.py | 30 ++++++++++++++++++++++++++++++
 tests/test_cli.py   | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 81 insertions(+)

diff --git a/docs/cli.rst b/docs/cli.rst
index 5eeb507..9c5344c 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -621,6 +621,19 @@ If you omit the other table and other column references ``sqlite-utils`` will at
 
 See :ref:`python_api_add_foreign_key` in the Python API documentation for further details, including how the automatic table guessing mechanism works.
 
+.. _cli_add_foreign_keys:
+
+Adding multiple foreign keys at once
+------------------------------------
+
+Adding a foreign key requires a ``VACUUM``. On large databases this can be an expensive operation, so if you are adding multiple foreign keys you can combine them into one operation (and hence one ``VACUUM``) using ``add-foreign-keys``::
+
+    $ sqlite-utils add-foreign-keys books.db \
+        books author_id authors id \
+        authors country_id countries id
+
+When you are using this command each foreign key needs to be defined in full, as four arguments - the table, column, other table and other column.
+
 .. _cli_index_foreign_keys:
 
 Adding indexes for all foreign keys
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 4a01cc8..eb44a8e 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -323,6 +323,36 @@ def add_foreign_key(path, table, column, other_table, other_column):
         raise click.ClickException(e)
 
 
+@cli.command(name="add-foreign-keys")
+@click.argument(
+    "path",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, allow_dash=False),
+    required=True,
+)
+@click.argument("foreign_key", nargs=-1)
+def add_foreign_keys(path, foreign_key):
+    """
+    Add multiple new foreign key constraints to a database. Example usage:
+
+    \b
+    sqlite-utils add-foreign-keys my.db \\
+        books author_id authors id \\
+        authors country_id countries id
+    """
+    db = sqlite_utils.Database(path)
+    if len(foreign_key) % 4 != 0:
+        raise click.ClickException(
+            "Each foreign key requires four values: table, column, other_table, other_column"
+        )
+    tuples = []
+    for i in range(len(foreign_key) // 4):
+        tuples.append(tuple(foreign_key[i * 4 : (i * 4) + 4]))
+    try:
+        db.add_foreign_keys(tuples)
+    except AlterError as e:
+        raise click.ClickException(e)
+
+
 @cli.command(name="index-foreign-keys")
 @click.argument(
     "path",
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 0a8c0ba..dba182b 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1405,3 +1405,41 @@ def test_query_update(db_path, args, expected):
     assert db.execute_returning_dicts("select * from dogs") == [
         {"id": 1, "age": 5, "name": "Cleo"},
     ]
+
+
+def test_add_foreign_keys(db_path):
+    db = Database(db_path)
+    db["countries"].insert({"id": 7, "name": "Panama"}, pk="id")
+    db["authors"].insert({"id": 3, "name": "Matilda", "country_id": 7}, pk="id")
+    db["books"].insert({"id": 2, "title": "Wolf anatomy", "author_id": 3}, pk="id")
+    assert db["authors"].foreign_keys == []
+    assert db["books"].foreign_keys == []
+    result = CliRunner().invoke(
+        cli.cli,
+        [
+            "add-foreign-keys",
+            db_path,
+            "authors",
+            "country_id",
+            "countries",
+            "id",
+            "books",
+            "author_id",
+            "authors",
+            "id",
+        ],
+    )
+    assert result.exit_code == 0
+    assert db["authors"].foreign_keys == [
+        ForeignKey(
+            table="authors",
+            column="country_id",
+            other_table="countries",
+            other_column="id",
+        )
+    ]
+    assert db["books"].foreign_keys == [
+        ForeignKey(
+            table="books", column="author_id", other_table="authors", other_column="id"
+        )
+    ]