add_foreign_key can now detect table and pk, refs #25

docs/cli.rst

@@ -300,12 +300,16 @@ Adding foreign key constraints
 
 The ``add-foreign-key`` command can be used to add new foreign key references to an existing table - something which SQLite's ``ALTER TABLE`` command does not support.
 
-See :ref:`python_api_add_foreign_key` in the Python API documentation for further details and warnings (this could corrupt your database).
-
 To add a foreign key constraint pointing the ``books.author_id`` column to ``authors.id`` in another table, do this::
 
     $ sqlite-utils add-foreign-key books.db books author_id authors id
 
+If you omit the other table and other column references ``sqlite-utils`` will attempt to guess them - so the above example could instead look like this::
+
+    $ sqlite-utils add-foreign-key books.db books author_id
+
+See :ref:`python_api_add_foreign_key` in the Python API documentation for further details, including how the automatic table guessing mechanism works.
+
 .. _cli_create_index:
 
 Creating indexes

docs/python-api.rst

@@ -334,6 +334,11 @@ Here's an example of this mechanism in action:
     ])
     db["books"].add_foreign_key("author_id", "authors", "id")
 
+The ``table.add_foreign_key(column, other_table, other_column)`` method takes the name of the column, the table that is being referenced and the key column within that other table. If you ommit the ``other_column`` argument the primary key from that table will be used automatically. If you omit the ``other_table`` argument the table will be guessed based on some simple rules:
+
+- 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_hash:
 
 Setting an ID based on the hash of the row contents

sqlite_utils/cli.py

@@ -200,8 +200,8 @@ def add_column(path, table, col_name, col_type, fk, fk_col, not_null_default):
 )
 @click.argument("table")
 @click.argument("column")
-@click.argument("other_table")
-@click.argument("other_column")
+@click.argument("other_table", required=False)
+@click.argument("other_column", required=False)
 def add_foreign_key(path, table, column, other_table, other_column):
     """
     Add a new foreign key constraint to an existing table. Example usage:

sqlite_utils/db.py

@@ -71,6 +71,14 @@ class AlterError(Exception):
     pass
 
 
+class NoObviousTable(Exception):
+    pass
+
+
+class BadPrimaryKey(Exception):
+    pass
+
+
 class Database:
     def __init__(self, filename_or_conn):
         if isinstance(filename_or_conn, str):
@@ -337,7 +345,43 @@ class Table:
     def drop(self):
         return self.db.conn.execute("DROP TABLE {}".format(self.name))
 
-    def add_foreign_key(self, column, other_table, other_column):
+    def guess_foreign_table(self, column):
+        column = column.lower()
+        possibilities = [column]
+        if column.endswith("_id"):
+            column_without_id = column[:-3]
+            possibilities.append(column_without_id)
+            if not column_without_id.endswith("s"):
+                possibilities.append(column_without_id + "s")
+        elif not column.endswith("s"):
+            possibilities.append(column + "s")
+        existing_tables = {t.lower(): t for t in self.db.table_names()}
+        for table in possibilities:
+            if table in existing_tables:
+                return existing_tables[table]
+        # If we get here there's no obvious candidate - raise an error
+        raise NoObviousTable(
+            "No obvious foreign key table for column '{}' - tried {}".format(
+                column, repr(possibilities)
+            )
+        )
+
+    def add_foreign_key(self, column, other_table=None, other_column=None):
+        # If other_table is not specified, attempt to guess it from the column
+        if other_table is None:
+            other_table = self.guess_foreign_table(column)
+        # If other_column is not specified, detect the primary key on other_table
+        if other_column is None:
+            pks = [c for c in self.db[other_table].columns if c.is_pk]
+            if len(pks) != 1:
+                raise BadPrimaryKey(
+                    "Could not detect single primary key for table '{}'".format(
+                        other_table
+                    )
+                )
+            else:
+                other_column = pks[0].name
+
         # Sanity check that the other column exists
         if (
             not [c for c in self.db[other_table].columns if c.name == other_column]

tests/test_cli.py

@@ -220,7 +220,18 @@ def test_add_column_not_null_default(db_path):
     )
 
 
-def test_add_foreign_key(db_path):
+@pytest.mark.parametrize(
+    "args,assert_message",
+    (
+        (
+            ["books", "author_id", "authors", "id"],
+            "Explicit other_table and other_column",
+        ),
+        (["books", "author_id", "authors"], "Explicit other_table, guess other_column"),
+        (["books", "author_id"], "Automatically guess other_table and other_column"),
+    ),
+)
+def test_add_foreign_key(db_path, args, assert_message):
     db = Database(db_path)
     db["authors"].insert_all(
         [{"id": 1, "name": "Sally"}, {"id": 2, "name": "Asheesh"}], pk="id"
@@ -232,13 +243,8 @@ def test_add_foreign_key(db_path):
         ]
     )
     assert (
-        0
-        == CliRunner()
-        .invoke(
-            cli.cli, ["add-foreign-key", db_path, "books", "author_id", "authors", "id"]
-        )
-        .exit_code
-    )
+        0 == CliRunner().invoke(cli.cli, ["add-foreign-key", db_path] + args).exit_code
+    ), assert_message
     assert [
         ForeignKey(
             table="books", column="author_id", other_table="authors", other_column="id"

tests/test_create.py

@@ -228,6 +228,17 @@ def test_add_column_foreign_key(fresh_db):
     )
 
 
+def test_add_foreign_key_guess_table(fresh_db):
+    fresh_db.create_table("dogs", {"name": str})
+    fresh_db.create_table("breeds", {"name": str, "id": int}, pk="id")
+    fresh_db["dogs"].add_column("breed_id", int)
+    fresh_db["dogs"].add_foreign_key("breed_id")
+    assert (
+        "CREATE TABLE [dogs] ( [name] TEXT , [breed_id] INTEGER, FOREIGN KEY(breed_id) REFERENCES breeds(id) )"
+        == collapse_whitespace(fresh_db["dogs"].schema)
+    )
+
+
 @pytest.mark.parametrize(
     "extra_data,expected_new_columns",
     [

tests/test_introspect.py

@@ -1,4 +1,5 @@
 from sqlite_utils.db import Index
+import pytest
 
 
 def test_table_names(existing_db):
@@ -79,3 +80,19 @@ def test_indexes(fresh_db):
         ),
         Index(seq=1, name="Gosh_c1", unique=0, origin="c", partial=0, columns=["c1"]),
     ] == fresh_db["Gosh"].indexes
+
+
+@pytest.mark.parametrize(
+    "column,expected_table_guess",
+    (
+        ("author", "authors"),
+        ("author_id", "authors"),
+        ("authors", "authors"),
+        ("genre", "genre"),
+        ("genre_id", "genre"),
+    ),
+)
+def test_guess_foreign_table(fresh_db, column, expected_table_guess):
+    fresh_db.create_table("authors", {"name": str})
+    fresh_db.create_table("genre", {"name": str})
+    assert expected_table_guess == fresh_db["books"].guess_foreign_table(column)