fix less than lookup on encrypted fields

Author: timgraham

django_mongodb_backend/lookups.py

@@ -4,6 +4,8 @@
     BuiltinLookup,
     FieldGetDbPrepValueIterableMixin,
     IsNull,
+    LessThan,
+    LessThanOrEqual,
     Lookup,
     PatternLookup,
     UUIDTextMixin,
@@ -101,6 +103,26 @@ def is_null_path(self, compiler, connection):
     return connection.mongo_operators["isnull"](lhs_mql, self.rhs)
 
 
+def less_than_path(self, compiler, connection):
+    lhs_mql = process_lhs(self, compiler, connection)
+    value = process_rhs(self, compiler, connection)
+    # Encrypted fields don't support null and Automatic Encryption cannot
+    # handle it ("csfle "analyze_query" failed: typenull type isn't supported
+    # for the range encrypted index.), so omit the null check.
+    if getattr(self.lhs.output_field, "encrypted", False):
+        return {lhs_mql: {"$lt": value}}
+    return connection.mongo_operators[self.lookup_name](lhs_mql, value)
+
+
+def less_than_or_equal_path(self, compiler, connection):
+    lhs_mql = process_lhs(self, compiler, connection)
+    value = process_rhs(self, compiler, connection)
+    # Same comment as less_than_path.
+    if getattr(self.lhs.output_field, "encrypted", False):
+        return {lhs_mql: {"$lte": value}}
+    return connection.mongo_operators[self.lookup_name](lhs_mql, value)
+
+
 # from https://www.pcre.org/current/doc/html/pcre2pattern.html#SEC4
 REGEX_MATCH_ESCAPE_CHARS = (
     ("\\", r"\\"),  # general escape character
@@ -157,6 +179,8 @@ def register_lookups():
     In.get_subquery_wrapping_pipeline = get_subquery_wrapping_pipeline
     IsNull.as_mql_expr = is_null_expr
     IsNull.as_mql_path = is_null_path
+    LessThan.as_mql_path = less_than_path
+    LessThanOrEqual.as_mql_path = less_than_or_equal_path
     Lookup.can_use_path = lookup_can_use_path
     PatternLookup.prep_lookup_value_mongo = pattern_lookup_prep_lookup_value
     UUIDTextMixin.as_mql = uuid_text_mixin

docs/ref/models/encrypted-fields.rst

@@ -65,12 +65,12 @@ The :ref:`available query types <manual:qe-fundamentals-encrypt-query>` depend
 on your version of MongoDB. For example, in MongoDB 8.0, the supported types
 are ``equality`` and ``range``.
 
-.. admonition:: Query types vs. Django lookups
+The supported lookups for ``equality`` queries are: :lookup:`exact` and
+lookup:`in`. The supported operators are AND (``&``) and OR (``|``).
 
-    Range queries in Queryable Encryption are different from Django's
-    :ref:`range lookups <django:field-lookups>`. Range queries allow you to
-    perform comparisons on encrypted fields, while Django's range lookups are
-    used for filtering based on a range of values.
+The supported lookups for ``range`` queries include those of ``equality``
+queries as well as :lookup:`lt`, :lookup:`lte`, :lookup:`gt`, and
+:lookup:`gte`.
 
 \* These fields don't support the ``queries`` argument:
 

tests/encryption_/test_fields.py

@@ -5,7 +5,7 @@
 
 from bson import ObjectId
 from django.db import DatabaseError
-from django.db.models import Avg, F
+from django.db.models import Avg, F, Q
 
 from django_mongodb_backend.fields import (
     EncryptedArrayField,
@@ -94,18 +94,41 @@ def test_array(self):
 
 class FieldTests(EncryptionTestCase):
     def assertEquality(self, model_cls, val):
-        model_cls.objects.create(value=val)
-        fetched = model_cls.objects.get(value=val)
-        self.assertEqual(fetched.value, val)
+        obj = model_cls.objects.create(value=val)
+        self.assertEqual(model_cls.objects.get(value=val), obj)
+        self.assertEqual(model_cls.objects.get(value__in=[val]), obj)
+        self.assertQuerySetEqual(model_cls.objects.exclude(value=val), [])
 
     def assertRange(self, model_cls, *, low, high, threshold):
-        model_cls.objects.create(value=low)
-        model_cls.objects.create(value=high)
+        obj1 = model_cls.objects.create(value=low)
+        obj2 = model_cls.objects.create(value=high)
         self.assertEqual(model_cls.objects.get(value=low).value, low)
         self.assertEqual(model_cls.objects.get(value=high).value, high)
-        objs = list(model_cls.objects.filter(value__gt=threshold))
-        self.assertEqual(len(objs), 1)
-        self.assertEqual(objs[0].value, high)
+        self.assertEqual(model_cls.objects.exclude(value=high).get().value, low)
+        self.assertCountEqual(model_cls.objects.filter(Q(value=high) | Q(value=low)), [obj1, obj2])
+        self.assertQuerySetEqual(
+            model_cls.objects.filter(value__gt=threshold), [high], attrgetter("value")
+        )
+        self.assertQuerySetEqual(
+            model_cls.objects.filter(value__gte=threshold), [high], attrgetter("value")
+        )
+        self.assertQuerySetEqual(
+            model_cls.objects.filter(value__lt=threshold), [low], attrgetter("value")
+        )
+        self.assertQuerySetEqual(
+            model_cls.objects.filter(value__lte=threshold), [low], attrgetter("value")
+        )
+        self.assertQuerySetEqual(
+            model_cls.objects.filter(value__in=[low]), [low], attrgetter("value")
+        )
+        msg = (
+            "Comparison disallowed between Queryable Encryption encrypted "
+            "fields and non-constant expressions; field 'value' is encrypted."
+        )
+        with self.assertRaisesMessage(DatabaseError, msg):
+            self.assertQuerySetEqual(
+                model_cls.objects.filter(value__lte=F("value")), [low], attrgetter("value")
+            )
 
     # Equality-only fields
     def test_binary(self):