rebase from main, update the test, this commit should be overwritten with the refactor

Author: timgraham

django_mongodb_backend/compiler.py

@@ -53,7 +53,7 @@ def _get_column_from_expression(self, expr, alias):
         Create a column named `alias` from the given expression to hold the
         aggregate value.
         """
-        column_target = expr.output_field.__class__()
+        column_target = expr.output_field.clone()
         column_target.db_column = alias
         column_target.set_attributes_from_name(alias)
         return Col(self.collection_name, column_target)
@@ -81,7 +81,7 @@ def _prepare_expressions_for_pipeline(self, expression, target, annotation_group
             alias = (
                 f"__aggregation{next(annotation_group_idx)}" if sub_expr != expression else target
             )
-            column_target = sub_expr.output_field.__class__()
+            column_target = sub_expr.output_field.clone()
             column_target.db_column = alias
             column_target.set_attributes_from_name(alias)
             inner_column = Col(self.collection_name, column_target)

django_mongodb_backend/fields/embedded_model.py

@@ -1,4 +1,5 @@
 from django.core import checks
+from django.core.exceptions import FieldDoesNotExist
 from django.db import models
 from django.db.models.fields.related import lazy_related_operation
 from django.db.models.lookups import Transform
@@ -112,7 +113,8 @@ def get_transform(self, name):
         transform = super().get_transform(name)
         if transform:
             return transform
-        return KeyTransformFactory(name)
+        field = self.embedded_model._meta.get_field(name)
+        return KeyTransformFactory(name, field)
 
     def validate(self, value, model_instance):
         super().validate(value, model_instance)
@@ -134,28 +136,65 @@ def formfield(self, **kwargs):
 
 
 class KeyTransform(Transform):
-    def __init__(self, key_name, *args, **kwargs):
+    def __init__(self, key_name, ref_field, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.key_name = str(key_name)
+        self.ref_field = ref_field
+
+    def get_transform(self, name):
+        result = None
+        if isinstance(self.ref_field, EmbeddedModelField):
+            opts = self.ref_field.embedded_model._meta
+            new_field = opts.get_field(name)
+            result = KeyTransformFactory(name, new_field)
+        else:
+            if self.ref_field.get_transform(name) is None:
+                raise FieldDoesNotExist(
+                    f"{self.ref_field.model._meta.object_name}.{self.ref_field.name}"
+                    f" has no field named '{name}'"
+                )
+            result = KeyTransformFactory(name, self.ref_field)
+        return result
 
     def preprocess_lhs(self, compiler, connection):
-        key_transforms = [self.key_name]
-        previous = self.lhs
+        previous = self
+        embedded_key_transforms = []
+        json_key_transforms = []
         while isinstance(previous, KeyTransform):
-            key_transforms.insert(0, previous.key_name)
+            if isinstance(previous.ref_field, EmbeddedModelField):
+                embedded_key_transforms.insert(0, previous.key_name)
+            else:
+                json_key_transforms.insert(0, previous.key_name)
             previous = previous.lhs
         mql = previous.as_mql(compiler, connection)
-        return mql, key_transforms
+        embedded_key_transforms.append(json_key_transforms.pop(0))
+        return mql, embedded_key_transforms, json_key_transforms
 
     def as_mql(self, compiler, connection):
-        mql, key_transforms = self.preprocess_lhs(compiler, connection)
+        mql, key_transforms, json_key_transforms = self.preprocess_lhs(compiler, connection)
         transforms = ".".join(key_transforms)
-        return f"{mql}.{transforms}"
+        result = f"{mql}.{transforms}"
+        for key in json_key_transforms:
+            get_field = {"$getField": {"input": result, "field": key}}
+            # Handle array indexing if the key is a digit. If key is something
+            # like '001', it's not an array index despite isdigit() returning True.
+            if key.isdigit() and str(int(key)) == key:
+                result = {
+                    "$cond": {
+                        "if": {"$isArray": result},
+                        "then": {"$arrayElemAt": [result, int(key)]},
+                        "else": get_field,
+                    }
+                }
+            else:
+                result = get_field
+        return result
 
 
 class KeyTransformFactory:
-    def __init__(self, key_name):
+    def __init__(self, key_name, ref_field):
         self.key_name = key_name
+        self.ref_field = ref_field
 
     def __call__(self, *args, **kwargs):
-        return KeyTransform(self.key_name, *args, **kwargs)
+        return KeyTransform(self.key_name, self.ref_field, *args, **kwargs)

tests/model_fields_/models.py

@@ -102,6 +102,7 @@ class Data(models.Model):
     integer = models.IntegerField(db_column="custom_column")
     auto_now = models.DateTimeField(auto_now=True)
     auto_now_add = models.DateTimeField(auto_now_add=True)
+    json_value = models.JSONField(default=dict)
 
 
 class Address(models.Model):
@@ -119,3 +120,10 @@ class Author(models.Model):
 class Book(models.Model):
     name = models.CharField(max_length=100)
     author = EmbeddedModelField(Author)
+
+
+class Library(models.Model):
+    name = models.CharField(max_length=100)
+    books = models.ManyToManyField("Book", related_name="libraries")
+    location = models.CharField(max_length=100, null=True, blank=True)
+    best_seller = models.CharField(max_length=100, null=True, blank=True)

tests/model_fields_/test_embedded_model.py

@@ -1,5 +1,17 @@
-from django.core.exceptions import ValidationError
+import operator
+
+from django.core.exceptions import FieldDoesNotExist, ValidationError
 from django.db import models
+from django.db.models import (
+    Exists,
+    ExpressionWrapper,
+    F,
+    IntegerField,
+    Max,
+    OuterRef,
+    Subquery,
+    Sum,
+)
 from django.test import SimpleTestCase, TestCase
 from django.test.utils import isolate_apps
 
@@ -11,6 +23,7 @@
     Book,
     Data,
     Holder,
+    Library,
 )
 
 
@@ -104,6 +117,85 @@ def test_nested(self):
         )
         self.assertCountEqual(Book.objects.filter(author__address__city="NYC"), [obj])
 
+    def test_nested_not_exists(self):
+        msg = "Address.city has no field named 'president'"
+        with self.assertRaisesMessage(FieldDoesNotExist, msg):
+            Book.objects.filter(author__address__city__president="NYC")
+
+    def test_not_exists_in_embedded(self):
+        msg = "Address has no field named 'floor'"
+        with self.assertRaisesMessage(FieldDoesNotExist, msg):
+            Book.objects.filter(author__address__floor="NYC")
+
+    def test_embedded_with_json_field(self):
+        models = []
+        for i in range(4):
+            m = Holder.objects.create(
+                data=Data(json_value={"field1": i * 5, "field2": {"0": {"value": list(range(i))}}})
+            )
+            models.append(m)
+
+        all_models = Holder.objects.all()
+
+        self.assertCountEqual(
+            Holder.objects.filter(data__json_value__field2__0__value__0=0),
+            models[1:],
+        )
+        self.assertCountEqual(
+            Holder.objects.filter(data__json_value__field2__0__value__1=1),
+            models[2:],
+        )
+        self.assertCountEqual(Holder.objects.filter(data__json_value__field2__0__value__1=5), [])
+
+        self.assertCountEqual(Holder.objects.filter(data__json_value__field1__lt=100), all_models)
+        self.assertCountEqual(Holder.objects.filter(data__json_value__field1__gt=100), [])
+        self.assertCountEqual(
+            Holder.objects.filter(
+                data__json_value__field1__gte=5, data__json_value__field1__lte=10
+            ),
+            models[1:3],
+        )
+
+    def truncate_ms(self, value):
+        """Truncate microsends to millisecond precision as supported by MongoDB."""
+        return value.replace(microsecond=(value.microsecond // 1000) * 1000)
+
+    def test_ordering_by_embedded_field(self):
+        query = Holder.objects.filter(data__integer__gt=3).order_by("-data__integer").values("pk")
+        expected = [{"pk": e.pk} for e in list(reversed(self.objs[4:]))]
+        self.assertSequenceEqual(query, expected)
+
+    def test_ordering_grouping_by_embedded_field(self):
+        expected = sorted(
+            (Holder.objects.create(data=Data(integer=x)) for x in range(6)),
+            key=lambda x: x.data.integer,
+        )
+        query = (
+            Holder.objects.annotate(
+                group=ExpressionWrapper(F("data__integer") + 5, output_field=IntegerField())
+            )
+            .values("group")
+            .annotate(max_auto_now=Max("data__auto_now"))
+            .order_by("data__integer")
+        )
+        query_response = [{**e, "max_auto_now": self.truncate_ms(e["max_auto_now"])} for e in query]
+        self.assertSequenceEqual(
+            query_response,
+            [
+                {"group": e.data.integer + 5, "max_auto_now": self.truncate_ms(e.data.auto_now)}
+                for e in expected
+            ],
+        )
+
+    def test_ordering_grouping_by_sum(self):
+        [Holder.objects.create(data=Data(integer=x)) for x in range(6)]
+        qs = (
+            Holder.objects.values("data__integer")
+            .annotate(sum=Sum("data__integer"))
+            .order_by("sum")
+        )
+        self.assertQuerySetEqual(qs, [0, 2, 4, 6, 8, 10], operator.itemgetter("sum"))
+
 
 @isolate_apps("model_fields_")
 class CheckTests(SimpleTestCase):
@@ -123,3 +215,62 @@ class MyModel(models.Model):
         self.assertEqual(
             msg, "Embedded models cannot have relational fields (Target.key is a ForeignKey)."
         )
+
+
+class SubqueryExistsTest(TestCase):
+    def setUp(self):
+        # Create test data
+        address1 = Address.objects.create(city="New York", state="NY", zip_code=10001)
+        address2 = Address.objects.create(city="Boston", state="MA", zip_code=20002)
+        author1 = Author.objects.create(name="Alice", age=30, address=address1)
+        author2 = Author.objects.create(name="Bob", age=40, address=address2)
+        book1 = Book.objects.create(name="Book A", author=author1)
+        book2 = Book.objects.create(name="Book B", author=author2)
+        Book.objects.create(name="Book C", author=author2)
+        Book.objects.create(name="Book D", author=author2)
+        Book.objects.create(name="Book E", author=author1)
+
+        library1 = Library.objects.create(
+            name="Central Library", location="Downtown", best_seller="Book A"
+        )
+        library2 = Library.objects.create(
+            name="Community Library", location="Suburbs", best_seller="Book A"
+        )
+
+        # Add books to libraries
+        library1.books.add(book1, book2)
+        library2.books.add(book2)
+
+    def test_exists_subquery(self):
+        subquery = Book.objects.filter(
+            author__name=OuterRef("name"), author__address__city="Boston"
+        )
+        queryset = Author.objects.filter(Exists(subquery))
+
+        self.assertEqual(queryset.count(), 1)
+
+    def test_in_subquery(self):
+        subquery = Author.objects.filter(age__gt=35).values("name")
+        queryset = Book.objects.filter(author__name__in=Subquery(subquery)).order_by("name")
+
+        self.assertEqual(queryset.count(), 3)
+        self.assertQuerySetEqual(queryset, ["Book B", "Book C", "Book D"], lambda book: book.name)
+
+    def test_range_query(self):
+        queryset = Author.objects.filter(age__range=(25, 45)).order_by("name")
+
+        self.assertEqual(queryset.count(), 2)
+        self.assertQuerySetEqual(queryset, ["Alice", "Bob"], lambda author: author.name)
+
+    def test_exists_with_foreign_object(self):
+        subquery = Library.objects.filter(best_seller=OuterRef("name"))
+        queryset = Book.objects.filter(Exists(subquery))
+
+        self.assertEqual(queryset.count(), 1)
+        self.assertEqual(queryset.first().name, "Book A")
+
+    def test_foreign_field_with_ranges(self):
+        queryset = Library.objects.filter(books__author__age__range=(25, 35))
+
+        self.assertEqual(queryset.count(), 1)
+        self.assertEqual(queryset.first().name, "Central Library")