diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index bbc6177..9f50a53 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1048,6 +1048,26 @@ class Table: self.create_index(column_values.keys(), unique=True) return pk + def m2m(self, other_table, record_or_list, pk=None): + our_id = self.last_pk + records = ( + [record_or_list] + if not isinstance(record_or_list, (list, tuple)) + else record_or_list + ) + tables = list(sorted([self.name, other_table])) + columns = ["{}_id".format(t) for t in tables] + m2m_table = self.db.table( + "{}_{}".format(*tables), pk=columns, foreign_keys=columns + ) + # Ensure each record exists in other table + for record in records: + id = self.db[other_table].upsert(record, pk=pk).last_pk + m2m_table.upsert( + {"{}_id".format(other_table): 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..e7d6e30 --- /dev/null +++ b/tests/test_m2m.py @@ -0,0 +1,43 @@ +from sqlite_utils.db import ForeignKey +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