From d645032cfa4edbccd0542eecdddca29edf9f7b07 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Wed, 12 Jun 2019 21:51:09 -0700 Subject: [PATCH] add_foreign_key can now detect table and pk, refs #25 --- docs/cli.rst | 8 +++++-- docs/python-api.rst | 5 +++++ sqlite_utils/cli.py | 4 ++-- sqlite_utils/db.py | 46 +++++++++++++++++++++++++++++++++++++++- tests/test_cli.py | 22 ++++++++++++------- tests/test_create.py | 11 ++++++++++ tests/test_introspect.py | 17 +++++++++++++++ 7 files changed, 100 insertions(+), 13 deletions(-) diff --git a/docs/cli.rst b/docs/cli.rst index 27994ba..f9ef84c 100644 --- a/docs/cli.rst +++ b/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 diff --git a/docs/python-api.rst b/docs/python-api.rst index 41f13d8..216b46d 100644 --- a/docs/python-api.rst +++ b/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 diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 4a32ddb..0424eda 100644 --- a/sqlite_utils/cli.py +++ b/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: diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index a77d9da..3dbae66 100644 --- a/sqlite_utils/db.py +++ b/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] diff --git a/tests/test_cli.py b/tests/test_cli.py index f164635..24bbc38 100644 --- a/tests/test_cli.py +++ b/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" diff --git a/tests/test_create.py b/tests/test_create.py index b03e925..5653b18 100644 --- a/tests/test_create.py +++ b/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", [ diff --git a/tests/test_introspect.py b/tests/test_introspect.py index 73bcd24..a49028c 100644 --- a/tests/test_introspect.py +++ b/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)