Fix BelongsToMany where_has with a query callback

fastapi_startkit/src/fastapi_startkit/masoniteorm/relationships/BelongsToMany.py

@@ -3,6 +3,7 @@
 
 from .BaseRelationship import BaseRelationship, ManyResult, TModel
 from ..collection import Collection
+from ..models.builder import QueryBuilder
 from ..models.pivot import Pivot
 
 
@@ -401,31 +402,24 @@ def joins(self, builder, clause=None):
         return result
 
     def query_where_exists(self, builder, callback, method="where_exists"):
-        query = self.get_builder()
-        pivot_table = self._table or self.get_pivot_table_name(query, builder)
-        table = self.get_builder().get_table_name()
-
-        getattr(builder, method)(
-            query.new()
-            .table(table)
-            .join(
-                f"{pivot_table}",
-                f"{table}.{self.other_owner_key}",
-                "=",
-                f"{pivot_table}.{self.foreign_key}",
+        related = self.related_exists_query(builder)
+        constrained = callback(related)
+        if isinstance(constrained, QueryBuilder) and constrained is not related:
+            raise TypeError(
+                "where_has callback must constrain and return the query it receives, not a different QueryBuilder"
             )
-            .where_column(
-                f"{pivot_table}.{self.local_key}",
-                f"{builder.get_table_name()}.{self.local_owner_key}",
-            )
-            .where_in(self.other_owner_key, callback(query.select(self.other_owner_key)))
-        )
+        if constrained is not None and constrained is not related:
+            related.where_in(f"{related.get_table_name()}.{self.other_owner_key}", constrained)
+        return getattr(builder, method)(related)
 
     def query_has(self, current_query_builder, method="where_exists"):
+        return getattr(current_query_builder, method)(self.related_exists_query(current_query_builder))
+
+    def related_exists_query(self, builder):
         query = self.get_builder()
-        pivot_table = self._table or self.get_pivot_table_name(query, current_query_builder)
-        table = self.get_builder().get_table_name()
-        return getattr(current_query_builder, method)(
+        pivot_table = self._table or self.get_pivot_table_name(query, builder)
+        table = query.get_table_name()
+        return (
             query.new()
             .table(table)
             .join(
@@ -436,7 +430,7 @@ def query_has(self, current_query_builder, method="where_exists"):
             )
             .where_column(
                 f"{pivot_table}.{self.local_key}",
-                f"{current_query_builder.get_table_name()}.{self.local_owner_key}",
+                f"{builder.get_table_name()}.{self.local_owner_key}",
             )
         )
 

fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_belongs_to_many.py

@@ -130,9 +130,63 @@ async def test_query_has_filters_stores_with_products(self):
         self.assertIn(self.store.id, store_ids)
         self.assertNotIn(store2.id, store_ids)
 
-    async def test_query_has_returns_builder(self):
-        # query_has should add a where_exists clause to the builder
+    async def test_query_has_adds_exists_clause(self):
         builder = Store.query()
-        result = Store.products.query_has(builder, method="where_exists")
-        # The builder has a where clause appended (we just verify no error is raised)
-        self.assertIsNotNone(builder)
+        Store.products.query_has(builder, method="where_exists")
+        self.assertEqual(
+            builder.to_sql(),
+            'SELECT * FROM "stores" WHERE EXISTS (SELECT * FROM "products" '
+            'INNER JOIN "product_store" ON "products"."id" = "product_store"."product_id" '
+            "WHERE product_store.store_id = stores.id)",
+        )
+
+
+class TestBelongsToManyWhereHas(TestCase):
+    async def asyncSetUp(self):
+        await super().asyncSetUp()
+        self.downtown = await Store.create({"name": "Downtown"})
+        self.uptown = await Store.create({"name": "Uptown"})
+        self.empty = await Store.create({"name": "Empty"})
+        self.widget = await Product.create({"name": "Widget"})
+        self.gadget = await Product.create({"name": "Gadget"})
+        await Store.products.attach(self.downtown, self.widget)
+        await Store.products.attach(self.uptown, self.gadget)
+
+    async def names(self, query):
+        return sorted(store.name for store in await query.get())
+
+    def scoped(self):
+        return Store.where_in("id", [self.downtown.id, self.uptown.id, self.empty.id])
+
+    async def test_where_has_without_callback(self):
+        self.assertEqual(await self.names(self.scoped().where_has("products")), ["Downtown", "Uptown"])
+
+    async def test_where_has_with_query_callback(self):
+        query = self.scoped().where_has("products", lambda q: q.where("name", "Widget"))
+        self.assertEqual(await self.names(query), ["Downtown"])
+
+    async def test_where_has_with_query_callback_matching_nothing(self):
+        query = self.scoped().where_has("products", lambda q: q.where("name", "Missing"))
+        self.assertEqual(await self.names(query), [])
+
+    async def test_where_has_with_callback_returning_none(self):
+        def only_gadgets(q):
+            q.where("name", "Gadget")
+
+        self.assertEqual(await self.names(self.scoped().where_has("products", only_gadgets)), ["Uptown"])
+
+    async def test_where_has_rejects_callback_returning_a_different_builder(self):
+        with self.assertRaisesRegex(TypeError, "must constrain and return the query it receives"):
+            Store.where_has("products", lambda q: Product.where("name", "Widget"))
+
+    async def test_where_has_with_id_list_callback(self):
+        query = self.scoped().where_has("products", lambda q: [self.gadget.id])
+        self.assertEqual(await self.names(query), ["Uptown"])
+
+    async def test_or_where_has_with_query_callback(self):
+        query = Store.where("name", "Empty").or_where_has("products", lambda q: q.where("name", "Gadget"))
+        self.assertEqual(await self.names(query), ["Empty", "Uptown"])
+
+    async def test_or_where_has_without_callback(self):
+        query = Store.where("name", "Empty").or_where_has("products")
+        self.assertEqual(await self.names(query), ["Downtown", "Empty", "Uptown"])