From e8f887ef4a0977243811b90bc2ce9aed9d2c206a Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 30 Jun 2019 16:50:54 -0700
Subject: [PATCH] sqlite-utils index-foreign-keys / db.index_foreign_keys()

Closes #33
---
 docs/cli.rst         |  9 +++++++++
 docs/python-api.rst  | 11 +++++++++++
 sqlite_utils/cli.py  | 14 ++++++++++++++
 sqlite_utils/db.py   | 10 ++++++++++
 tests/test_cli.py    | 10 ++++++++++
 tests/test_create.py | 10 ++++++++++
 6 files changed, 64 insertions(+)

diff --git a/docs/cli.rst b/docs/cli.rst
index 8e47ad6..db1954c 100644
--- a/docs/cli.rst
+++ b/docs/cli.rst
@@ -310,6 +310,15 @@ 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_index_foreign_keys:
+
+Adding indexes for all foreign keys
+-----------------------------------
+
+If you want to ensure that every foreign key column in your database has a corresponding index, you can do so like this::
+
+    $ sqlite-utils index-foreign-keys books.db
+
 .. _cli_defaults_not_null:
 
 Setting defaults and not null constraints
diff --git a/docs/python-api.rst b/docs/python-api.rst
index 8928c03..c7ef65a 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -419,6 +419,17 @@ Here's an example adding two foreign keys at once:
         ("dogs", "home_town_id", "towns", "id")
     ])
 
+.. _python_api_index_foreign_keys:
+
+Adding indexes for all foreign keys
+-----------------------------------
+
+If you want to ensure that every foreign key column in your database has a corresponding index, you can do so like this:
+
+.. code-block:: python
+
+    db.index_foreign_keys()
+
 .. _python_api_hash:
 
 Setting an ID based on the hash of the row contents
diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py
index 0df540a..1c8661f 100644
--- a/sqlite_utils/cli.py
+++ b/sqlite_utils/cli.py
@@ -217,6 +217,20 @@ def add_foreign_key(path, table, column, other_table, other_column):
         raise click.ClickException(e)
 
 
+@cli.command(name="index-foreign-keys")
+@click.argument(
+    "path",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, allow_dash=False),
+    required=True,
+)
+def index_foreign_keys(path):
+    """
+    Ensure every foreign key column has an index on it.
+    """
+    db = sqlite_utils.Database(path)
+    db.index_foreign_keys()
+
+
 @cli.command(name="create-index")
 @click.argument(
     "path",
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 582017f..3a44cb4 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -315,6 +315,16 @@ class Database:
         # can see the newly created foreign key.
         self.vacuum()
 
+    def index_foreign_keys(self):
+        for table_name in self.table_names():
+            table = self[table_name]
+            existing_indexes = {
+                i.columns[0] for i in table.indexes if len(i.columns) == 1
+            }
+            for fk in table.foreign_keys:
+                if fk.column not in existing_indexes:
+                    table.create_index([fk.column])
+
     def vacuum(self):
         self.conn.execute("VACUUM;")
 
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 0316a51..47afbb0 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -310,6 +310,16 @@ def test_add_column_foreign_key(db_path):
     assert "table 'bobcats' does not exist" in str(result.exception)
 
 
+def test_index_foreign_keys(db_path):
+    test_add_column_foreign_key(db_path)
+    db = Database(db_path)
+    assert [] == db["books"].indexes
+    result = CliRunner().invoke(cli.cli, ["index-foreign-keys", db_path])
+    assert [["author_id"], ["author_name_ref"]] == [
+        i.columns for i in db["books"].indexes
+    ]
+
+
 def test_enable_fts(db_path):
     assert None == Database(db_path)["Gosh"].detect_fts()
     result = CliRunner().invoke(
diff --git a/tests/test_create.py b/tests/test_create.py
index 88d7676..5e76d0c 100644
--- a/tests/test_create.py
+++ b/tests/test_create.py
@@ -361,6 +361,16 @@ def test_add_foreign_key_guess_table(fresh_db):
     )
 
 
+def test_index_foreign_keys(fresh_db):
+    test_add_foreign_key_guess_table(fresh_db)
+    assert [] == fresh_db["dogs"].indexes
+    fresh_db.index_foreign_keys()
+    assert [["breed_id"]] == [i.columns for i in fresh_db["dogs"].indexes]
+    # Calling it a second time should do nothing
+    fresh_db.index_foreign_keys()
+    assert [["breed_id"]] == [i.columns for i in fresh_db["dogs"].indexes]
+
+
 @pytest.mark.parametrize(
     "extra_data,expected_new_columns",
     [