Add combinable test

Author: WaVEV

django_mongodb_backend/compiler.py

@@ -17,7 +17,7 @@
 from django.utils.functional import cached_property
 from pymongo import ASCENDING, DESCENDING
 
-from .expressions.builtins import SearchExpression
+from .expressions.builtins import SearchExpression, SearchVector
 from .query import MongoQuery, wrap_database_errors
 
 
@@ -109,37 +109,26 @@ def _prepare_expressions_for_pipeline(self, expression, target, annotation_group
             replacements[sub_expr] = self._get_replace_expr(sub_expr, group, alias)
         return replacements, group
 
-    def _prepare_search_expressions_for_pipeline(
-        self, expression, target, search_idx, replacements
-    ):
+    def _prepare_search_expressions_for_pipeline(self, expression, search_idx, replacements):
         searches = {}
         for sub_expr in self._get_search_expressions(expression):
             if sub_expr not in replacements:
                 alias = f"__search_expr.search{next(search_idx)}"
                 replacements[sub_expr] = self._get_replace_expr(sub_expr, searches, alias)
-        return list(searches.values())
 
     def _prepare_search_query_for_aggregation_pipeline(self, order_by):
         replacements = {}
-        searches = []
         annotation_group_idx = itertools.count(start=1)
-        for target, expr in self.query.annotation_select.items():
-            expr_searches = self._prepare_search_expressions_for_pipeline(
-                expr, target, annotation_group_idx, replacements
-            )
-            searches += expr_searches
+        for expr in self.query.annotation_select.values():
+            self._prepare_search_expressions_for_pipeline(expr, annotation_group_idx, replacements)
 
         for expr, _ in order_by:
-            expr_searches = self._prepare_search_expressions_for_pipeline(
-                expr, None, annotation_group_idx, replacements
-            )
-            searches += expr_searches
+            self._prepare_search_expressions_for_pipeline(expr, annotation_group_idx, replacements)
 
-        having_group = self._prepare_search_expressions_for_pipeline(
-            self.having, None, annotation_group_idx, replacements
+        self._prepare_search_expressions_for_pipeline(
+            self.having, annotation_group_idx, replacements
         )
-        searches += having_group
-        return searches, replacements
+        return replacements
 
     def _prepare_annotations_for_aggregation_pipeline(self, order_by):
         """Prepare annotations for the aggregation pipeline."""
@@ -248,22 +237,36 @@ def _build_aggregation_pipeline(self, ids, group):
                 pipeline.append({"$unset": "_id"})
         return pipeline
 
-    def _compound_searches_queries(self, searches, search_replacements):
-        if not searches:
+    def _compound_searches_queries(self, search_replacements):
+        if not search_replacements:
             return []
-        if len(searches) > 1:
+        if len(search_replacements) > 1:
             raise ValueError("Cannot perform more than one search operation.")
-        score_function = "searchScore" if "$search" in searches[0] else "vectorSearchScore"
-        return [searches[0], {"$addFields": {"__search_expr.search1": {"$meta": score_function}}}]
+        pipeline = []
+        for search, result_col in search_replacements.items():
+            score_function = (
+                "vectorSearchScore" if isinstance(search, SearchVector) else "searchScore"
+            )
+            pipeline.extend(
+                [
+                    search.as_mql(self, self.connection),
+                    {
+                        "$addFields": {
+                            result_col.as_mql(self, self.connection).removeprefix("$"): {
+                                "$meta": score_function
+                            }
+                        }
+                    },
+                ]
+            )
+        return pipeline
 
     def pre_sql_setup(self, with_col_aliases=False):
         extra_select, order_by, group_by = super().pre_sql_setup(with_col_aliases=with_col_aliases)
-        searches, search_replacements = self._prepare_search_query_for_aggregation_pipeline(
-            order_by
-        )
+        search_replacements = self._prepare_search_query_for_aggregation_pipeline(order_by)
         group, group_replacements = self._prepare_annotations_for_aggregation_pipeline(order_by)
         all_replacements = {**search_replacements, **group_replacements}
-        self.search_pipeline = self._compound_searches_queries(searches, search_replacements)
+        self.search_pipeline = self._compound_searches_queries(search_replacements)
         # query.group_by is either:
         # - None: no GROUP BY
         # - True: group by select fields

django_mongodb_backend/creation.py

@@ -22,6 +22,9 @@ def _destroy_test_db(self, test_database_name, verbosity):
 
         for collection in self.connection.introspection.table_names():
             if not collection.startswith("system."):
+                db_collection = self.connection.database.get_collection(collection)
+                for search_indexes in db_collection.list_search_indexes():
+                    db_collection.drop_search_index(search_indexes["name"])
                 self.connection.database.drop_collection(collection)
 
     def create_test_db(self, *args, **kwargs):

tests/expressions_/test_combinable_search_expression.py

@@ -0,0 +1,76 @@
+from django.test import SimpleTestCase
+
+from django_mongodb_backend.expressions.builtins import (
+    CombinedSearchExpression,
+    CompoundExpression,
+    SearchEquals,
+)
+
+
+class CombinedSearchExpressionResolutionTest(SimpleTestCase):
+    def test_combined_expression_and_or_not_resolution(self):
+        A = SearchEquals(path="headline", value="A")
+        B = SearchEquals(path="headline", value="B")
+        C = SearchEquals(path="headline", value="C")
+        D = SearchEquals(path="headline", value="D")
+        expr = (~A | B) & (C | D)
+        solved = CombinedSearchExpression.resolve(expr)
+        self.assertIsInstance(solved, CompoundExpression)
+        solved_A = CompoundExpression(must_not=[CompoundExpression(must=[A])])
+        solved_B = CompoundExpression(must=[B])
+        solved_C = CompoundExpression(must=[C])
+        solved_D = CompoundExpression(must=[D])
+        self.assertCountEqual(solved.must[0].should, [solved_A, solved_B])
+        self.assertEqual(solved.must[0].minimum_should_match, 1)
+        self.assertEqual(solved.must[1].should, [solved_C, solved_D])
+
+    def test_combined_expression_de_morgans_resolution(self):
+        A = SearchEquals(path="headline", value="A")
+        B = SearchEquals(path="headline", value="B")
+        C = SearchEquals(path="headline", value="C")
+        D = SearchEquals(path="headline", value="D")
+        expr = ~(A | B) & (C | D)
+        solved_A = CompoundExpression(must_not=[CompoundExpression(must=[A])])
+        solved_B = CompoundExpression(must_not=[CompoundExpression(must=[B])])
+        solved_C = CompoundExpression(must=[C])
+        solved_D = CompoundExpression(must=[D])
+        solved = CombinedSearchExpression.resolve(expr)
+        self.assertIsInstance(solved, CompoundExpression)
+        self.assertCountEqual(solved.must[0].must, [solved_A, solved_B])
+        self.assertEqual(solved.must[0].minimum_should_match, None)
+        self.assertEqual(solved.must[1].should, [solved_C, solved_D])
+        self.assertEqual(solved.minimum_should_match, None)
+
+    def test_combined_expression_doble_negation(self):
+        A = SearchEquals(path="headline", value="A")
+        expr = ~~A
+        solved = CombinedSearchExpression.resolve(expr)
+        solved_A = CompoundExpression(must=[A])
+        self.assertIsInstance(solved, CompoundExpression)
+        self.assertEqual(solved, solved_A)
+
+    def test_combined_expression_long_right_tree(self):
+        A = SearchEquals(path="headline", value="A")
+        B = SearchEquals(path="headline", value="B")
+        C = SearchEquals(path="headline", value="C")
+        D = SearchEquals(path="headline", value="D")
+        solved_A = CompoundExpression(must=[A])
+        solved_B = CompoundExpression(must_not=[CompoundExpression(must=[B])])
+        solved_C = CompoundExpression(must=[C])
+        solved_D = CompoundExpression(must=[D])
+        expr = A & ~(B & ~(C & D))
+        solved = CombinedSearchExpression.resolve(expr)
+        self.assertIsInstance(solved, CompoundExpression)
+        self.assertEqual(len(solved.must), 2)
+        self.assertEqual(solved.must[0], solved_A)
+        self.assertEqual(len(solved.must[1].should), 2)
+        self.assertEqual(solved.must[1].should[0], solved_B)
+        self.assertCountEqual(solved.must[1].should[1].must, [solved_C, solved_D])
+        expr = A | ~(B | ~(C | D))
+        solved = CombinedSearchExpression.resolve(expr)
+        self.assertIsInstance(solved, CompoundExpression)
+        self.assertEqual(len(solved.should), 2)
+        self.assertEqual(solved.should[0], solved_A)
+        self.assertEqual(len(solved.should[1].must), 2)
+        self.assertEqual(solved.should[1].must[0], solved_B)
+        self.assertCountEqual(solved.should[1].must[1].should, [solved_C, solved_D])

tests/queries_/models.py

@@ -60,4 +60,4 @@ class Article(models.Model):
     number = models.IntegerField()
     body = models.TextField()
     location = models.JSONField(null=True)
-    plot_embedding = ArrayField(models.FloatField(), size=3)
+    plot_embedding = ArrayField(models.FloatField(), size=3, null=True)