From e8d958109ee290cfa1b44ef7a39629bb50ab673e Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sun, 14 Nov 2021 14:49:19 -0800
Subject: [PATCH] create_index(..., find_unique_name=True) option, refs #335

---
 docs/python-api.rst  |  4 ++-
 sqlite_utils/db.py   | 59 +++++++++++++++++++++++++++++++-------------
 tests/test_create.py | 17 +++++++++++++
 3 files changed, 62 insertions(+), 18 deletions(-)

diff --git a/docs/python-api.rst b/docs/python-api.rst
index 20d080d..8932656 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -2141,7 +2141,9 @@ You can create an index on a table using the ``.create_index(columns)`` method.
 
     db["dogs"].create_index(["is_good_dog"])
 
-By default the index will be named ``idx_{table-name}_{columns}`` - if you want to customize the name of the created index you can pass the ``index_name`` parameter:
+By default the index will be named ``idx_{table-name}_{columns}``. If you pass ``find_unique_name=True`` and the automatically derived name already exists, an available name will be found by incrementing a suffix number, for example ``idx_items_title_2``.
+
+You can customize the name of the created index by passing the ``index_name`` parameter:
 
 .. code-block:: python
 
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index ab7252a..1e75861 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -899,7 +899,7 @@ class Database:
             }
             for fk in table.foreign_keys:
                 if fk.column not in existing_indexes:
-                    table.create_index([fk.column])
+                    table.create_index([fk.column], find_unique_name=True)
 
     def vacuum(self):
         "Run a SQLite ``VACUUM`` against the database."
@@ -1521,6 +1521,7 @@ class Table(Queryable):
         index_name: Optional[str] = None,
         unique: bool = False,
         if_not_exists: bool = False,
+        find_unique_name: bool = False,
     ):
         """
         Create an index on this table.
@@ -1530,6 +1531,8 @@ class Table(Queryable):
         - ``index_name`` - the name to use for the new index. Defaults to the column names joined on ``_``.
         - ``unique`` - should the index be marked as unique, forcing unique values?
         - ``if_not_exists`` - only create the index if one with that name does not already exist.
+        - ``find_unique_name`` - if ``index_name`` is not provided and the automatically derived name
+          already exists, keep incrementing a suffix number to find an available name.
 
         See :ref:`python_api_create_index`.
         """
@@ -1544,23 +1547,45 @@ class Table(Queryable):
             else:
                 fmt = "[{}]"
             columns_sql.append(fmt.format(column))
-        sql = (
-            textwrap.dedent(
-                """
-            CREATE {unique}INDEX {if_not_exists}[{index_name}]
-                ON [{table_name}] ({columns});
-        """
+
+        suffix = None
+        while True:
+            sql = (
+                textwrap.dedent(
+                    """
+                CREATE {unique}INDEX {if_not_exists}[{index_name}]
+                    ON [{table_name}] ({columns});
+            """
+                )
+                .strip()
+                .format(
+                    index_name="{}_{}".format(index_name, suffix)
+                    if suffix
+                    else index_name,
+                    table_name=self.name,
+                    columns=", ".join(columns_sql),
+                    unique="UNIQUE " if unique else "",
+                    if_not_exists="IF NOT EXISTS " if if_not_exists else "",
+                )
             )
-            .strip()
-            .format(
-                index_name=index_name,
-                table_name=self.name,
-                columns=", ".join(columns_sql),
-                unique="UNIQUE " if unique else "",
-                if_not_exists="IF NOT EXISTS " if if_not_exists else "",
-            )
-        )
-        self.db.execute(sql)
+            try:
+                self.db.execute(sql)
+                break
+            except OperationalError as e:
+                # find_unique_name=True - try again if 'index ... already exists'
+                arg = e.args[0]
+                if (
+                    find_unique_name
+                    and arg.startswith("index ")
+                    and arg.endswith(" already exists")
+                ):
+                    if suffix is None:
+                        suffix = 2
+                    else:
+                        suffix += 1
+                    continue
+                else:
+                    raise e
         return self
 
     def add_column(
diff --git a/tests/test_create.py b/tests/test_create.py
index 58dda95..8d5b023 100644
--- a/tests/test_create.py
+++ b/tests/test_create.py
@@ -4,6 +4,7 @@ from sqlite_utils.db import (
     DescIndex,
     AlterError,
     NoObviousTable,
+    OperationalError,
     ForeignKey,
     Table,
     View,
@@ -752,6 +753,22 @@ def test_create_index_desc(fresh_db):
     )
 
 
+def test_create_index_find_unique_name():
+    db = Database(memory=True)
+    table = db["t"]
+    table.insert({"id": 1})
+    table.create_index(["id"])
+    # Without find_unique_name should error
+    with pytest.raises(OperationalError, match="index idx_t_id already exists"):
+        table.create_index(["id"])
+    # With find_unique_name=True it should work
+    table.create_index(["id"], find_unique_name=True)
+    table.create_index(["id"], find_unique_name=True)
+    # Should have three now
+    index_names = {idx.name for idx in table.indexes}
+    assert index_names == {"idx_t_id", "idx_t_id_2", "idx_t_id_3"}
+
+
 @pytest.mark.parametrize(
     "data_structure",
     (