Add get_subquery_wrapping_pipeline and unit test

Author: WaVEV

django_mongodb_backend/fields/embedded_model_array.py

@@ -129,7 +129,40 @@ def as_mql(self, compiler, connection):
 
 @_EmbeddedModelArrayOutputField.register_lookup
 class EmbeddedModelArrayFieldIn(EmbeddedModelArrayFieldBuiltinLookup, lookups.In):
-    pass
+    def get_subquery_wrapping_pipeline(self, compiler, connection, field_name, expr):
+        return [
+            {
+                "$facet": {
+                    "group": [
+                        {"$project": {"tmp_name": expr.as_mql(compiler, connection)}},
+                        {
+                            "$unwind": "$tmp_name",
+                        },
+                        {
+                            "$group": {
+                                "_id": None,
+                                "tmp_name": {"$addToSet": "$tmp_name"},
+                            }
+                        },
+                    ]
+                }
+            },
+            {
+                "$project": {
+                    field_name: {
+                        "$ifNull": [
+                            {
+                                "$getField": {
+                                    "input": {"$arrayElemAt": ["$group", 0]},
+                                    "field": "tmp_name",
+                                }
+                            },
+                            [],
+                        ]
+                    }
+                }
+            },
+        ]
 
 
 @_EmbeddedModelArrayOutputField.register_lookup
@@ -171,6 +204,41 @@ class EmbeddedModelArrayFieldAll(EmbeddedModelArrayFieldBuiltinLookup, Lookup):
     lookup_name = "all"
     get_db_prep_lookup_value_is_iterable = False
 
+    def get_subquery_wrapping_pipeline(self, compiler, connection, field_name, expr):
+        return [
+            {
+                "$facet": {
+                    "group": [
+                        {"$project": {"tmp_name": expr.as_mql(compiler, connection)}},
+                        {
+                            "$unwind": "$tmp_name",
+                        },
+                        {
+                            "$group": {
+                                "_id": None,
+                                "tmp_name": {"$addToSet": "$tmp_name"},
+                            }
+                        },
+                    ]
+                }
+            },
+            {
+                "$project": {
+                    field_name: {
+                        "$ifNull": [
+                            {
+                                "$getField": {
+                                    "input": {"$arrayElemAt": ["$group", 0]},
+                                    "field": "tmp_name",
+                                }
+                            },
+                            [],
+                        ]
+                    }
+                }
+            },
+        ]
+
     def as_mql(self, compiler, connection):
         lhs_mql = process_lhs(self, compiler, connection)
         values = process_rhs(self, compiler, connection)

tests/model_fields_/models.py

@@ -180,6 +180,11 @@ class ArtifactDetail(EmbeddedModel):
     last_restoration = EmbeddedModelField(RestorationRecord, null=True)
 
 
+class ExhibitAudit(models.Model):
+    related_section_number = models.IntegerField()
+    reviewed = models.BooleanField()
+
+
 # A section within an exhibit, containing multiple artifacts.
 class ExhibitSection(EmbeddedModel):
     section_number = models.IntegerField()

tests/model_fields_/test_embedded_model_array.py

@@ -10,6 +10,7 @@
 
 from .models import (
     ArtifactDetail,
+    ExhibitAudit,
     ExhibitSection,
     Movie,
     MuseumExhibit,
@@ -124,6 +125,7 @@ def setUpTestData(cls):
                     ],
                 )
             ],
+            main_section=ExhibitSection(section_number=2),
         )
         cls.lost_empires = MuseumExhibit.objects.create(
             exhibit_name="Lost Empires",
@@ -154,6 +156,9 @@ def setUpTestData(cls):
         cls.egypt_tour = Tour.objects.create(guide="Amira", exhibit=cls.egypt)
         cls.wonders_tour = Tour.objects.create(guide="Carlos", exhibit=cls.wonders)
         cls.lost_tour = Tour.objects.create(guide="Yelena", exhibit=cls.lost_empires)
+        cls.audit_1 = ExhibitAudit.objects.create(related_section_number=1, reviewed=True)
+        cls.audit_2 = ExhibitAudit.objects.create(related_section_number=2, reviewed=True)
+        cls.audit_3 = ExhibitAudit.objects.create(related_section_number=5, reviewed=False)
 
     def test_filter_with_field(self):
         self.assertCountEqual(
@@ -309,6 +314,35 @@ def test_foreign_field_with_slice(self):
         qs = Tour.objects.filter(exhibit__sections__0_2__section_number__all=[1, 2])
         self.assertEqual(list(qs), [self.wonders_tour])
 
+    def test_subquery_section_number_lt(self):
+        subq = ExhibitAudit.objects.filter(
+            related_section_number__in=models.OuterRef("sections__section_number")
+        ).values("related_section_number")[:1]
+        self.assertCountEqual(
+            MuseumExhibit.objects.filter(sections__section_number=subq),
+            [self.egypt, self.wonders, self.new_descoveries],
+        )
+
+    def test_check_all_subquery(self):
+        subquery = ExhibitAudit.objects.filter(reviewed=True).values_list(
+            "related_section_number", flat=True
+        )
+        result = MuseumExhibit.objects.filter(sections__section_number__all=subquery)
+        self.assertCountEqual(result, [self.wonders])
+
+    def test_check_in_subquery(self):
+        subquery = ExhibitAudit.objects.filter(reviewed=True).values_list(
+            "related_section_number", flat=True
+        )
+        result = MuseumExhibit.objects.filter(sections__section_number__in=subquery)
+        self.assertCountEqual(result, [self.wonders, self.egypt, self.new_descoveries])
+
+    def test_array_as_rhs(self):
+        result = MuseumExhibit.objects.filter(
+            main_section__section_number__in=models.F("sections__section_number")
+        )
+        self.assertCountEqual(result, [self.new_descoveries])
+
 
 @isolate_apps("model_fields_")
 class CheckTests(SimpleTestCase):