foreign_keys=[...] can now detect tables and keys - closes #25

This commit is contained in:
Simon Willison 2019-06-12 22:32:26 -07:00
commit 07a10bdb2c
3 changed files with 125 additions and 25 deletions

View file

@ -149,7 +149,25 @@ Any operation that can create a table (``.create()``, ``.insert()``, ``.insert_a
If you are using your database with `Datasette <https://datasette.readthedocs.io/>`__, 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:

View file

@ -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 (

View file

@ -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"},