diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index f0c29b5..1c8661f 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -228,13 +228,7 @@ def index_foreign_keys(path): Ensure every foreign key column has an index on it. """ db = sqlite_utils.Database(path) - for table_name in db.table_names(): - table = db[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: - print("Creating index on {}.{}".format(table_name, fk.column)) - table.create_index([fk.column]) + db.index_foreign_keys() @cli.command(name="create-index") 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_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", [