From ba1211d4456911bf0bd13f2e753a56ed988df3b4 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 3 Aug 2019 17:28:03 +0300
Subject: [PATCH] Implemented .m2m(table, lookup=...)

---
 sqlite_utils/db.py | 28 +++++++++++++++++++---------
 tests/test_m2m.py  | 31 +++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 9f50a53..e759910 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -1048,21 +1048,31 @@ class Table:
             self.create_index(column_values.keys(), unique=True)
             return pk
 
-    def m2m(self, other_table, record_or_list, pk=None):
+    def m2m(self, other_table, record_or_list=None, pk=None, lookup=None):
         our_id = self.last_pk
-        records = (
-            [record_or_list]
-            if not isinstance(record_or_list, (list, tuple))
-            else record_or_list
-        )
+        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]))
         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
+        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 = self.db[other_table].upsert(record, pk=pk).last_pk
+                m2m_table.upsert(
+                    {"{}_id".format(other_table): id, "{}_id".format(self.name): our_id}
+                )
+        else:
+            id = self.db[other_table].lookup(lookup)
             m2m_table.upsert(
                 {"{}_id".format(other_table): id, "{}_id".format(self.name): our_id}
             )
diff --git a/tests/test_m2m.py b/tests/test_m2m.py
index c7e593d..471166e 100644
--- a/tests/test_m2m.py
+++ b/tests/test_m2m.py
@@ -43,6 +43,37 @@ def test_insert_m2m_list(fresh_db):
     ] == dogs_humans.foreign_keys
 
 
+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_argument(fresh_db):
     # .m2m("humans", ..., m2m_table="relationships")
     assert False