Refactor and fix unit tests.

Author: WaVEV

django_mongodb/base.py

@@ -119,15 +119,6 @@ def _isnull_operator(a, b):
         "iregex": lambda a, b: regex_match(a, b, insensitive=True),
     }
 
-    mongo_operators_idx = {
-        "exact": lambda a, b: {a: {"$eq": b}},
-        "gt": lambda a, b: {a: {"$gt": b}},
-        "gte": lambda a, b: {a: {"$gte": b}},
-        "lt": lambda a, b: {a: {"$lt": b}},
-        "lte": lambda a, b: {a: {"$lte": b}},
-        "in": lambda a, b: {a: {"$in": b}},
-    }
-
     display_name = "MongoDB"
     vendor = "mongodb"
     Database = Database

django_mongodb/indexes.py

@@ -2,12 +2,22 @@
 from django.db import NotSupportedError
 from django.db.models import Index
 from django.db.models.expressions import Col
+from django.db.models.fields.related_lookups import In
 from django.db.models.lookups import BuiltinLookup
 from django.db.models.sql.query import Query
 from django.db.models.sql.where import AND, XOR, WhereNode
 
 from .query_utils import process_rhs
 
+mongo_operators_idx = {
+    "exact": lambda a, b: {a: {"$eq": b}},
+    "gt": lambda a, b: {a: {"$gt": b}},
+    "gte": lambda a, b: {a: {"$gte": b}},
+    "lt": lambda a, b: {a: {"$lt": b}},
+    "lte": lambda a, b: {a: {"$lte": b}},
+    "in": lambda a, b: {a: {"$in": b}},
+}
+
 
 def _get_condition_mql(self, model, schema_editor):
     """Analogous to Index._get_condition_sql()."""
@@ -23,13 +33,21 @@ def builtin_lookup_idx(self, compiler, connection):
     lhs_mql = self.lhs.target.column
     value = process_rhs(self, compiler, connection)
     try:
-        return connection.mongo_operators_idx[self.lookup_name](lhs_mql, value)
+        return mongo_operators_idx[self.lookup_name](lhs_mql, value)
     except KeyError:
         raise NotSupportedError(
-            f"MongoDB does not support the {self.lookup_name} lookup in indexes."
+            f"MongoDB does not support the '{self.lookup_name}' lookup in indexes."
         ) from None
 
 
+def in_idx(self, compiler, connection):
+    if not connection.features.is_mongodb_6_0:
+        raise NotSupportedError(
+            f"MongoDB does not support the {self.lookup_name} lookup in indexes."
+        )
+    return builtin_lookup_idx(self, compiler, connection)
+
+
 def where_node_idx(self, compiler, connection):
     if self.connector == AND:
         full_needed, empty_needed = len(self.children), 1
@@ -38,13 +56,13 @@ def where_node_idx(self, compiler, connection):
     if self.connector == AND:
         operator = "$and"
     elif self.connector == XOR:
-        raise NotSupportedError("Xor in indexes is not supported.")
+        raise NotSupportedError("MongoDB does not support the 'Xor' lookup in indexes.")
     else:
         if not connection.features.is_mongodb_6_0:
-            raise NotSupportedError("Or in indexes is not supported.")
+            raise NotSupportedError("MongoDB does not support the 'Or' lookup in indexes.")
         operator = "$or"
     if self.negated:
-        raise NotSupportedError("Negated field in indexes is not supported.")
+        raise NotSupportedError("MongoDB does not support negated field in indexes.")
     children_mql = []
     for child in self.children:
         try:
@@ -77,5 +95,6 @@ def where_node_idx(self, compiler, connection):
 
 def register_indexes():
     BuiltinLookup.as_mql_idx = builtin_lookup_idx
+    In.as_mql_idx = in_idx
     Index._get_condition_mql = _get_condition_mql
     WhereNode.as_mql_idx = where_node_idx

tests/indexes_/models.py

@@ -7,6 +7,3 @@ class Article(models.Model):
     published = models.BooleanField(default=False)
     number = models.IntegerField()
     text = models.CharField(max_length=100)
-
-    class Meta:
-        indexes = [models.Index(fields=["headline", "pub_date"])]

tests/indexes_/test_mql.py

@@ -9,29 +9,16 @@
 
 class PartialIndexTests(TestCase):
     def test_not_supported(self):
-        with connection.schema_editor() as editor:
-            index = Index(
+        msg = "MongoDB does not support the 'isnull' lookup in indexes."
+        with connection.schema_editor() as editor, self.assertRaisesMessage(NotSupportedError, msg):
+            Index(
                 name="not_supported",
                 fields=["headline"],
                 condition=Q(pk__isnull=True) & Q(pk__contains="test1"),
-            )
-            self.assertEqual(
-                index._get_condition_mql(Article, schema_editor=editor),
-                {},
-            )
-            editor.add_index(index=index, model=Article)
-            with connection.cursor() as cursor:
-                self.assertIn(
-                    index.name,
-                    connection.introspection.get_constraints(
-                        cursor=cursor,
-                        table_name=Article._meta.db_table,
-                    ),
-                )
-            editor.remove_index(index=index, model=Article)
+            )._get_condition_mql(Article, schema_editor=editor)
 
     def test_negated_not_supported(self):
-        msg = "Negated field in indexes is not supported."
+        msg = "MongoDB does not support negated field in indexes."
         with self.assertRaisesMessage(NotSupportedError, msg), connection.schema_editor() as editor:
             Index(
                 name="test",
@@ -40,7 +27,7 @@ def test_negated_not_supported(self):
             )._get_condition_mql(Article, schema_editor=editor)
 
     def test_xor_not_supported(self):
-        msg = "Xor in indexes is not supported."
+        msg = "MongoDB does not support the 'Xor' lookup in indexes."
         with self.assertRaisesMessage(NotSupportedError, msg), connection.schema_editor() as editor:
             Index(
                 name="test",
@@ -50,12 +37,22 @@ def test_xor_not_supported(self):
 
     @skipIfDBFeature("is_mongodb_6_0")
     def test_or_not_supported(self):
-        msg = "Or in indexes is not supported."
+        msg = "MongoDB does not support the 'Or' lookup in indexes."
         with self.assertRaisesMessage(NotSupportedError, msg), connection.schema_editor() as editor:
             Index(
                 name="test",
                 fields=["headline"],
-                condition=Q(pk=True) ^ Q(pk=False),
+                condition=Q(pk=True) | Q(pk=False),
+            )._get_condition_mql(Article, schema_editor=editor)
+
+    @skipIfDBFeature("is_mongodb_6_0")
+    def test_in_not_supported(self):
+        msg = "MongoDB does not support the 'in' lookup in indexes."
+        with self.assertRaisesMessage(NotSupportedError, msg), connection.schema_editor() as editor:
+            Index(
+                name="test",
+                fields=["headline"],
+                condition=Q(pk__in=[True]),
             )._get_condition_mql(Article, schema_editor=editor)
 
     @skipUnlessDBFeature("is_mongodb_6_0")