diff --git a/docs/python-api.rst b/docs/python-api.rst index 216b46d..fa7b691 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -149,7 +149,25 @@ Any operation that can create a table (``.create()``, ``.insert()``, ``.insert_a If you are using your database with `Datasette `__, Datasette will detect these constraints and use them to generate hyperlinks to associated records. -The ``foreign_keys`` argument takes a sequence of three-tuples, each one specifying the column, other table and other column that should be used to create the relationship. For example: +The ``foreign_keys`` argument takes a list that indicates which foreign keys should be created. The list can take several forms. The simplest is a list of columns: + +.. code-block:: python + + foreign_keys=["author_id"] + +The library will guess which tables you wish to reference based on the column names using the rules described in :ref:`python_api_add_foreign_key`. + +You can also be more explicit, by passing in a list of tuples: + +.. code-block:: python + + foreign_keys=[ + ("author_id", "authors", "id") + ] + +This means that the ``author_id`` column should be a foreign key that references the ``id`` column in the ``authors`` table. + +You can leave off the third item in the tuple to have the referenced column automatically set to the primary key of that table. A full example: .. code-block:: python @@ -161,7 +179,7 @@ The ``foreign_keys`` argument takes a sequence of three-tuples, each one specify {"title": "Hedgehogs of the world", "author_id": 1}, {"title": "How to train your wolf", "author_id": 2}, ], foreign_keys=[ - ("author_id", "authors", "id") + ("author_id", "authors") ]) .. _python_api_bulk_inserts: diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 3dbae66..e223782 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -112,11 +112,52 @@ class Database: keys = [d[0] for d in cursor.description] return [dict(zip(keys, row)) for row in cursor.fetchall()] + def resolve_foreign_keys(self, name, foreign_keys): + # foreign_keys may be a list of strcolumn names, a list of ForeignKey tuples, + # a list of tuple-pairs or a list of tuple-triples. We want to turn + # it into a list of ForeignKey tuples + if all(isinstance(fk, ForeignKey) for fk in foreign_keys): + return foreign_keys + if all(isinstance(fk, str) for fk in foreign_keys): + # It's a list of columns + fks = [] + for column in foreign_keys: + other_table = self[name].guess_foreign_table(column) + other_column = self[name].guess_foreign_column(other_table) + fks.append(ForeignKey(name, column, other_table, other_column)) + return fks + assert all( + isinstance(fk, (tuple, list)) for fk in foreign_keys + ), "foreign_keys= should be a list of tuples" + fks = [] + for tuple_or_list in foreign_keys: + assert len(tuple_or_list) in ( + 2, + 3, + ), "foreign_keys= should be a list of tuple pairs or triples" + if len(tuple_or_list) == 3: + fks.append( + ForeignKey( + name, tuple_or_list[0], tuple_or_list[1], tuple_or_list[2] + ) + ) + else: + # Guess the primary key + fks.append( + ForeignKey( + name, + tuple_or_list[0], + tuple_or_list[1], + self[name].guess_foreign_column(tuple_or_list[1]), + ) + ) + return fks + def create_table( self, name, columns, pk=None, foreign_keys=None, column_order=None, hash_id=None ): - foreign_keys = foreign_keys or [] - foreign_keys_by_name = {fk[0]: fk for fk in foreign_keys} + foreign_keys = self.resolve_foreign_keys(name, foreign_keys or []) + foreign_keys_by_column = {fk.column: fk for fk in foreign_keys} column_items = list(columns.items()) if column_order is not None: column_items.sort( @@ -126,12 +167,12 @@ class Database: column_items.insert(0, (hash_id, str)) pk = hash_id # Sanity check foreign_keys point to existing tables - for _, fk_other_table, fk_other_column in foreign_keys: + for fk in foreign_keys: if not any( - c for c in self[fk_other_table].columns if c.name == fk_other_column + c for c in self[fk.other_table].columns if c.name == fk.other_column ): raise AlterError( - "No such column: {}.{}".format(fk_other_table, fk_other_column) + "No such column: {}.{}".format(fk.other_table, fk.other_column) ) extra = "" columns_sql = ",\n".join( @@ -141,10 +182,10 @@ class Database: primary_key=" PRIMARY KEY" if (pk == col_name) else "", references=( " REFERENCES [{other_table}]([{other_column}])".format( - other_table=foreign_keys_by_name[col_name][1], - other_column=foreign_keys_by_name[col_name][2], + other_table=foreign_keys_by_column[col_name].other_table, + other_column=foreign_keys_by_column[col_name].other_column, ) - if col_name in foreign_keys_by_name + if col_name in foreign_keys_by_column else "" ), ) @@ -366,21 +407,22 @@ class Table: ) ) + def guess_foreign_column(self, other_table): + 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: + return pks[0].name + 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 + other_column = self.guess_foreign_column(other_table) # Sanity check that the other column exists if ( diff --git a/tests/test_create.py b/tests/test_create.py index 5653b18..36a4374 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -1,4 +1,11 @@ -from sqlite_utils.db import Index, Database, ForeignKey, AlterError +from sqlite_utils.db import ( + Index, + Database, + ForeignKey, + AlterError, + NoObviousTable, + ForeignKey, +) import collections import datetime import json @@ -91,13 +98,46 @@ def test_create_table_column_order(fresh_db): ] == [{"name": col.name, "type": col.type} for col in fresh_db["table"].columns] -def test_create_table_works_for_m2m_with_only_foreign_keys(fresh_db): +@pytest.mark.parametrize( + "foreign_key_specification,expected_exception", + ( + # You can specify triples, pairs, or a list of columns + ((("one_id", "one", "id"), ("two_id", "two", "id")), False), + ((("one_id", "one"), ("two_id", "two")), False), + (("one_id", "two_id"), False), + # You can also specify ForeignKey tuples: + ( + ( + ForeignKey("m2m", "one_id", "one", "id"), + ForeignKey("m2m", "two_id", "two", "id"), + ), + False, + ), + # If you specify a column that doesn't point to a table, you get an error: + (("one_id", "two_id", "three_id"), NoObviousTable), + # Tuples of the wrong length get an error: + ((("one_id", "one", "id", "five"), ("two_id", "two", "id")), AssertionError), + # Likewise a bad column: + ((("one_id", "one", "id2"),), AlterError), + # Or a list of dicts + (({"one_id": "one"},), AssertionError), + ), +) +def test_create_table_works_for_m2m_with_only_foreign_keys( + fresh_db, foreign_key_specification, expected_exception +): fresh_db["one"].insert({"id": 1}, pk="id") fresh_db["two"].insert({"id": 1}, pk="id") - fresh_db["m2m"].insert( - {"one_id": 1, "two_id": 1}, - foreign_keys=(("one_id", "one", "id"), ("two_id", "two", "id")), - ) + if expected_exception: + with pytest.raises(expected_exception): + fresh_db["m2m"].insert( + {"one_id": 1, "two_id": 1}, foreign_keys=foreign_key_specification + ) + return + else: + fresh_db["m2m"].insert( + {"one_id": 1, "two_id": 1}, foreign_keys=foreign_key_specification + ) assert [ {"name": "one_id", "type": "INTEGER"}, {"name": "two_id", "type": "INTEGER"},