table.m2m(...) method, with tests

This commit is contained in:
Simon Willison 2019-07-31 08:31:27 +03:00
commit 35eeafaaa3
2 changed files with 63 additions and 0 deletions

View file

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

43
tests/test_m2m.py Normal file
View file

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