Add subquery support for EmbeddedModelArrayField

Author: WaVEV

django_mongodb_backend/fields/embedded_model_array.py

@@ -143,7 +143,50 @@ 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):
+ # This pipeline is adapted from that of ArrayField, because the
+ # structure of EmbeddedModelArrayField on the RHS behaves similar to
+ # ArrayField.
+ return [
+ {
+ "$facet": {
+ "gathered_data": [
+ {"$project": {"tmp_name": expr.as_mql(compiler, connection)}},
+ # To concatenate all the values from the RHS subquery,
+ # use an $unwind followed by a $group.
+ {
+ "$unwind": "$tmp_name",
+ },
+ # The $group stage collects values into an array using
+ # $addToSet. The use of {_id: null} results in a
+ # single grouped array. However, because arrays from
+ # multiple documents are aggregated, the result is a
+ # list of lists.
+ {
+ "$group": {
+ "_id": None,
+ "tmp_name": {"$addToSet": "$tmp_name"},
+ }
+ },
+ ]
+ }
+ },
+ {
+ "$project": {
+ field_name: {
+ "$ifNull": [
+ {
+ "$getField": {
+ "input": {"$arrayElemAt": ["$gathered_data", 0]},
+ "field": "tmp_name",
+ }
+ },
+ [],
+ ]
+ }
+ }
+ },
+ ]
 
 
 @_EmbeddedModelArrayOutputField.register_lookup

docs/source/releases/5.2.x.rst

@@ -2,6 +2,16 @@
 Django MongoDB Backend 5.2.x
 ============================
 
+5.2.0 beta 2
+============
+
+*Unreleased*
+
+New features
+------------
+
+- Added subquery support for :class:`~.fields.EmbeddedModelArrayField`.
+
 5.2.0 beta 1
 ============
 

tests/model_fields_/models.py

@@ -206,6 +206,15 @@ def __str__(self):
  return f"Restored by {self.restored_by} on {self.date}"
 
 
+# An audit of a section in the museum.
+class Audit(models.Model):
+ section_number = models.IntegerField()
+ reviewed = models.BooleanField()
+
+ def __str__(self):
+ return f"Section {self.section_number} audit"
+
+
 # ForeignKey to a model with EmbeddedModelArrayField.
 class Tour(models.Model):
  guide = models.CharField(max_length=100)

tests/model_fields_/test_embedded_model_array.py

@@ -1,14 +1,17 @@
+import unittest
 from datetime import date
+from operator import attrgetter
 
 from django.core.exceptions import FieldDoesNotExist
 from django.db import connection, models
+from django.db.models.expressions import Value
 from django.test import SimpleTestCase, TestCase
 from django.test.utils import CaptureQueriesContext, isolate_apps
 
-from django_mongodb_backend.fields import EmbeddedModelArrayField
+from django_mongodb_backend.fields import ArrayField, EmbeddedModelArrayField
 from django_mongodb_backend.models import EmbeddedModel
 
-from .models import Artifact, Exhibit, Movie, Restoration, Review, Section, Tour
+from .models import Artifact, Audit, Exhibit, Movie, Restoration, Review, Section, Tour
 
 
 class MethodTests(SimpleTestCase):
@@ -116,6 +119,7 @@ def setUpTestData(cls):
  ],
  )
  ],
+ main_section=Section(number=2),
  )
  cls.lost_empires = Exhibit.objects.create(
  name="Lost Empires",
@@ -146,6 +150,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 = Audit.objects.create(section_number=1, reviewed=True)
+ cls.audit_2 = Audit.objects.create(section_number=2, reviewed=True)
+ cls.audit_3 = Audit.objects.create(section_number=5, reviewed=False)
 
  def test_exact(self):
  self.assertCountEqual(
@@ -284,6 +291,71 @@ def test_foreign_field_with_slice(self):
  qs = Tour.objects.filter(exhibit__sections__0_2__number__in=[1, 2])
  self.assertCountEqual(qs, [self.wonders_tour, self.egypt_tour])
 
+ def test_subquery_numeric_lookups(self):
+ subquery = Audit.objects.filter(
+ section_number__in=models.OuterRef("sections__number")
+ ).values("section_number")[:1]
+ tests = [
+ ("exact", [self.egypt, self.new_descoveries, self.wonders]),
+ ("lt", []),
+ ("lte", [self.egypt, self.new_descoveries, self.wonders]),
+ ("gt", [self.wonders]),
+ ("gte", [self.egypt, self.new_descoveries, self.wonders]),
+ ]
+ for lookup, expected in tests:
+ with self.subTest(lookup=lookup):
+ kwargs = {f"sections__number__{lookup}": subquery}
+ self.assertCountEqual(Exhibit.objects.filter(**kwargs), expected)
+
+ def test_subquery_in_lookup(self):
+ subquery = Audit.objects.filter(reviewed=True).values_list("section_number", flat=True)
+ result = Exhibit.objects.filter(sections__number__in=subquery)
+ self.assertCountEqual(result, [self.wonders, self.new_descoveries, self.egypt])
+
+ def test_array_as_rhs(self):
+ result = Exhibit.objects.filter(main_section__number__in=models.F("sections__number"))
+ self.assertCountEqual(result, [self.new_descoveries])
+
+ def test_array_annotation_lookup(self):
+ result = Exhibit.objects.annotate(section_numbers=models.F("main_section__number")).filter(
+ section_numbers__in=models.F("sections__number")
+ )
+ self.assertCountEqual(result, [self.new_descoveries])
+
+ def test_array_as_rhs_for_arrayfield_lookups(self):
+ tests = [
+ ("exact", [self.wonders]),
+ ("lt", [self.new_descoveries]),
+ ("lte", [self.wonders, self.new_descoveries]),
+ ("gt", [self.egypt, self.lost_empires]),
+ ("gte", [self.egypt, self.wonders, self.lost_empires]),
+ ("overlap", [self.egypt, self.wonders, self.new_descoveries]),
+ ("contained_by", [self.wonders]),
+ ("contains", [self.egypt, self.wonders, self.new_descoveries, self.lost_empires]),
+ ]
+ for lookup, expected in tests:
+ with self.subTest(lookup=lookup):
+ kwargs = {f"section_numbers__{lookup}": models.F("sections__number")}
+ result = Exhibit.objects.annotate(
+ section_numbers=Value(
+ [1, 2], output_field=ArrayField(base_field=models.IntegerField())
+ )
+ ).filter(**kwargs)
+ self.assertCountEqual(result, expected)
+
+ @unittest.expectedFailure
+ def test_array_annotation_index(self):
+ # Slicing and indexing over an annotated EmbeddedModelArrayField would
+ # require a refactor of annotation handling.
+ result = Exhibit.objects.annotate(section_numbers=models.F("sections__number")).filter(
+ section_numbers__0=1
+ )
+ self.assertCountEqual(result, [self.new_descoveries, self.egypt])
+
+ def test_array_annotation(self):
+ qs = Exhibit.objects.annotate(section_numbers=models.F("sections__number")).order_by("name")
+ self.assertQuerySetEqual(qs, [[1], [], [2], [1, 2]], attrgetter("section_numbers"))
+
 
 @isolate_apps("model_fields_")
 class CheckTests(SimpleTestCase):