From b9256413d26875c2bc3841e68b90d3842e88ccb8 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Sat, 3 Aug 2019 21:07:06 +0300
Subject: [PATCH] db.m2m_table_candidates(table, other_table)

---
 sqlite_utils/db.py | 11 +++++++++++
 tests/test_m2m.py  | 22 ++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index fdc4a99..f52b2e8 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(
diff --git a/tests/test_m2m.py b/tests/test_m2m.py
index 10a3d50..5cfe29f 100644
--- a/tests/test_m2m.py
+++ b/tests/test_m2m.py
@@ -84,6 +84,28 @@ def test_m2m_explicit_table_name_argument(fresh_db):
     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 toble with fks to both tables
     # and use that if it exists.