diff --git a/docs/python-api.rst b/docs/python-api.rst index b1ad43d..de3c076 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -442,6 +442,95 @@ To extract the ``species`` column out to a separate ``Species`` table, you can d "species": "Common Juniper" }, extracts={"species": "Species"}) +.. _python_api_m2m: + +Working with many-to-many relationships +======================================= + +``sqlite-utils`` includes a shortcut for creating records using many-to-many relationships in the form of the ``table.m2m(...)`` method. + +Here's how to create two new records and connect them via a many-to-many table in a single line of code: + +.. code-block:: python + + db["dogs"].insert({"id": 1, "name": "Cleo"}, pk="id").m2m( + "humans", {"id": 1, "name": "Natalie"}, pk="id" + ) + +Running this example actually creates three tables: ``dogs``, ``humans`` and a many-to-many ``dogs_humans`` table. It will insert a record into each of those tables. + +The ``.m2m()`` method executes against the last record that was affected by ``.insert()`` or ``.update()`` - the record identified by the ``table.last_pk`` property. To execute ``.m2m()`` against a specific record you can first select it by passing its primary key to ``.update()``: + +.. code-block:: python + + db["dogs"].update(1).m2m( + "humans", {"id": 2, "name": "Simon"}, pk="id" + ) + +The first argument to ``.m2m()`` can be either the name of a table as a string or it can be the table object itself. + +The second argument can be a single dictionary record or a list of dictionaries. Thesee dictionaries will be passed to ``.upsert()`` against the specified table. + +Here's alternative code that creates the dog record and adds two people to it: + +.. code-block:: python + + db = Database(memory=True) + dogs = db.table("dogs", pk="id") + humans = db.table("humans", pk="id") + dogs.insert({"id": 1, "name": "Cleo"}).m2m( + humans, [ + {"id": 1, "name": "Natalie"}, + {"id": 2, "name": "Simon"} + ] + ) + +The method will attempt to find an existing many-to-many table by looking for a table that has foreign key relationships against both of the tables in the relationship. + +If it cannot find such a table, it will create a new one using the names of the two tables - ``dogs_humans`` in this example. You can customize the name of this table using the ``m2m_table=`` argument to ``.m2m()``. + +It it finds multiple candidate tables with foreign keys to both of the specified tables it will raise a ``sqlite_utils.db.NoObviousTable`` exception. You can avoid this error by specifying the correct table using ``m2m_table=``. + +.. _python_api_m2m_lookup: + +Using m2m and lookup tables together +------------------------------------ + +You can work with (or create) lookup tables as part of a call to ``.m2m()`` using the ``lookup=`` parameter. This accepts the same argument as ``table.lookup()`` does - a dictionary of values that should be used to lookup or create a row in the lookup table. + +This example creates a dogs table, populates it, creates a characteristics table, populates that and sets up a many-to-many relationship between the two. It chains ``.m2m()`` twice to create two associated characteristics: + +.. code-block:: python + + db = Database(memory=True) + dogs = db.table("dogs", pk="id") + dogs.insert({"id": 1, "name": "Cleo"}).m2m( + "characteristics", lookup={ + "name": "Playful" + } + ).m2m( + "characteristics", lookup={ + "name": "Opinionated" + } + ) + +You can inspect the database to see the results like this:: + + >>> db.table_names() + ['dogs', 'characteristics', 'characteristics_dogs'] + >>> list(db["dogs"].rows) + [{'id': 1, 'name': 'Cleo'}] + >>> list(db["characteristics"].rows) + [{'id': 1, 'name': 'Playful'}, {'id': 2, 'name': 'Opinionated'}] + >>> list(db["characteristics_dogs"].rows) + [{'characteristics_id': 1, 'dogs_id': 1}, {'characteristics_id': 2, 'dogs_id': 1}] + >>> print(db["characteristics_dogs"].schema) + CREATE TABLE [characteristics_dogs] ( + [characteristics_id] INTEGER REFERENCES [characteristics]([id]), + [dogs_id] INTEGER REFERENCES [dogs]([id]), + PRIMARY KEY ([characteristics_id], [dogs_id]) + ) + .. _python_api_add_column: Adding columns diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index bbc6177..3fcb133 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -297,6 +297,17 @@ class Database: ) ) + def m2m_table_candidates(self, table, other_table): + "Returns potential m2m tables for arguments, based on FKs" + candidates = [] + tables = {table, other_table} + for table in self.tables: + # Does it have foreign keys to both table and other_table? + has_fks_to = {fk.other_table for fk in table.foreign_keys} + if has_fks_to.issuperset(tables): + candidates.append(table.name) + return candidates + def add_foreign_keys(self, foreign_keys): # foreign_keys is a list of explicit 4-tuples assert all( @@ -1048,6 +1059,60 @@ class Table: self.create_index(column_values.keys(), unique=True) return pk + def m2m( + self, other_table, record_or_list=None, pk=DEFAULT, lookup=None, m2m_table=None + ): + if isinstance(other_table, str): + other_table = self.db.table(other_table, pk=pk) + our_id = self.last_pk + if lookup is not None: + assert record_or_list is None, "Provide lookup= or record, not both" + else: + assert record_or_list is not None, "Provide lookup= or record, not both" + tables = list(sorted([self.name, other_table.name])) + columns = ["{}_id".format(t) for t in tables] + if m2m_table is not None: + m2m_table_name = m2m_table + else: + # Detect if there is a single, unambiguous option + candidates = self.db.m2m_table_candidates(self.name, other_table.name) + if len(candidates) == 1: + m2m_table_name = candidates[0] + elif len(candidates) > 1: + raise NoObviousTable( + "No single obvious m2m table for {}, {} - use m2m_table= parameter".format( + self.name, other_table.name + ) + ) + else: + # If not, create a new table + m2m_table_name = m2m_table or "{}_{}".format(*tables) + m2m_table = self.db.table(m2m_table_name, pk=columns, foreign_keys=columns) + if lookup is None: + records = ( + [record_or_list] + if not isinstance(record_or_list, (list, tuple)) + else record_or_list + ) + # Ensure each record exists in other table + for record in records: + id = other_table.upsert(record, pk=pk).last_pk + m2m_table.upsert( + { + "{}_id".format(other_table.name): id, + "{}_id".format(self.name): our_id, + } + ) + else: + id = other_table.lookup(lookup) + m2m_table.upsert( + { + "{}_id".format(other_table.name): id, + "{}_id".format(self.name): our_id, + } + ) + return self + def chunks(sequence, size): iterator = iter(sequence) diff --git a/tests/test_m2m.py b/tests/test_m2m.py new file mode 100644 index 0000000..b5b897b --- /dev/null +++ b/tests/test_m2m.py @@ -0,0 +1,156 @@ +from sqlite_utils.db import ForeignKey, NoObviousTable +import pytest + + +def test_insert_m2m_single(fresh_db): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo"}, pk="id").m2m( + "humans", {"id": 1, "name": "Natalie D"}, pk="id" + ) + assert {"dogs_humans", "humans", "dogs"} == set(fresh_db.table_names()) + humans = fresh_db["humans"] + dogs_humans = fresh_db["dogs_humans"] + assert [{"id": 1, "name": "Natalie D"}] == list(humans.rows) + assert [{"humans_id": 1, "dogs_id": 1}] == list(dogs_humans.rows) + + +def test_insert_m2m_list(fresh_db): + dogs = fresh_db["dogs"] + dogs.insert({"id": 1, "name": "Cleo"}, pk="id").m2m( + "humans", + [{"id": 1, "name": "Natalie D"}, {"id": 2, "name": "Simon W"}], + pk="id", + ) + assert {"dogs", "humans", "dogs_humans"} == set(fresh_db.table_names()) + humans = fresh_db["humans"] + dogs_humans = fresh_db["dogs_humans"] + assert [{"humans_id": 1, "dogs_id": 1}, {"humans_id": 2, "dogs_id": 1}] == list( + dogs_humans.rows + ) + assert [{"id": 1, "name": "Natalie D"}, {"id": 2, "name": "Simon W"}] == list( + humans.rows + ) + assert [ + ForeignKey( + table="dogs_humans", column="dogs_id", other_table="dogs", other_column="id" + ), + ForeignKey( + table="dogs_humans", + column="humans_id", + other_table="humans", + other_column="id", + ), + ] == dogs_humans.foreign_keys + + +def test_m2m_with_table_objects(fresh_db): + dogs = fresh_db.table("dogs", pk="id") + humans = fresh_db.table("humans", pk="id") + dogs.insert({"id": 1, "name": "Cleo"}).m2m( + humans, [{"id": 1, "name": "Natalie D"}, {"id": 2, "name": "Simon W"}] + ) + expected_tables = {"dogs", "humans", "dogs_humans"} + assert expected_tables == set(fresh_db.table_names()) + assert 1 == dogs.count + assert 2 == humans.count + assert 2 == fresh_db["dogs_humans"].count + + +def test_m2m_lookup(fresh_db): + people = fresh_db.table("people", pk="id") + people.insert({"name": "Wahyu"}).m2m("tags", lookup={"tag": "Coworker"}) + people_tags = fresh_db["people_tags"] + tags = fresh_db["tags"] + assert people_tags.exists + assert tags.exists + assert [ + ForeignKey( + table="people_tags", + column="people_id", + other_table="people", + other_column="id", + ), + ForeignKey( + table="people_tags", column="tags_id", other_table="tags", other_column="id" + ), + ] == people_tags.foreign_keys + assert [{"people_id": 1, "tags_id": 1}] == list(people_tags.rows) + assert [{"id": 1, "name": "Wahyu"}] == list(people.rows) + assert [{"id": 1, "tag": "Coworker"}] == list(tags.rows) + + +def test_m2m_requires_either_records_or_lookup(fresh_db): + people = fresh_db.table("people", pk="id").insert({"name": "Wahyu"}) + with pytest.raises(AssertionError): + people.m2m("tags") + with pytest.raises(AssertionError): + people.m2m("tags", {"tag": "hello"}, lookup={"foo": "bar"}) + + +def test_m2m_explicit_table_name_argument(fresh_db): + people = fresh_db.table("people", pk="id") + people.insert({"name": "Wahyu"}).m2m( + "tags", lookup={"tag": "Coworker"}, m2m_table="tagged" + ) + assert fresh_db["tags"].exists + assert fresh_db["tagged"].exists + assert not fresh_db["people_tags"].exists + + +def test_m2m_table_candidates(fresh_db): + fresh_db.create_table("one", {"id": int, "name": str}, pk="id") + fresh_db.create_table("two", {"id": int, "name": str}, pk="id") + fresh_db.create_table("three", {"id": int, "name": str}, pk="id") + # No candidates at first + assert [] == fresh_db.m2m_table_candidates("one", "two") + # Create a candidate + fresh_db.create_table( + "one_m2m_two", {"one_id": int, "two_id": int}, foreign_keys=["one_id", "two_id"] + ) + assert ["one_m2m_two"] == fresh_db.m2m_table_candidates("one", "two") + # Add another table and there should be two candidates + fresh_db.create_table( + "one_m2m_two_and_three", + {"one_id": int, "two_id": int, "three_id": int}, + foreign_keys=["one_id", "two_id", "three_id"], + ) + assert {"one_m2m_two", "one_m2m_two_and_three"} == set( + fresh_db.m2m_table_candidates("one", "two") + ) + + +def test_uses_existing_m2m_table_if_exists(fresh_db): + # Code should look for an existing table with fks to both tables + # and use that if it exists. + people = fresh_db.create_table("people", {"id": int, "name": str}, pk="id") + fresh_db["tags"].lookup({"tag": "Coworker"}) + fresh_db.create_table( + "tagged", + {"people_id": int, "tags_id": int}, + foreign_keys=["people_id", "tags_id"], + ) + people.insert({"name": "Wahyu"}).m2m("tags", lookup={"tag": "Coworker"}) + assert fresh_db["tags"].exists + assert fresh_db["tagged"].exists + assert not fresh_db["people_tags"].exists + assert not fresh_db["tags_people"].exists + assert [{"people_id": 1, "tags_id": 1}] == list(fresh_db["tagged"].rows) + + +def test_requires_explicit_m2m_table_if_multiple_options(fresh_db): + # If the code scans for m2m tables and finds more than one candidate + # it should require that the m2m_table=x argument is used + people = fresh_db.create_table("people", {"id": int, "name": str}, pk="id") + fresh_db["tags"].lookup({"tag": "Coworker"}) + fresh_db.create_table( + "tagged", + {"people_id": int, "tags_id": int}, + foreign_keys=["people_id", "tags_id"], + ) + fresh_db.create_table( + "tagged2", + {"people_id": int, "tags_id": int}, + foreign_keys=["people_id", "tags_id"], + ) + with pytest.raises(NoObviousTable): + people.insert({"name": "Wahyu"}).m2m("tags", lookup={"tag": "Coworker"})