Add SearchExpressions

Author: WaVEV

.pre-commit-config.yaml

@@ -81,4 +81,4 @@ repos:
   rev: "v2.2.6"
   hooks:
   - id: codespell
-    args: ["-L", "nin"]
+    args: ["-L", "nin", "SearchIn", "searchin"]

django_mongodb_backend/compiler.py

@@ -17,7 +17,7 @@
 from django.utils.functional import cached_property
 from pymongo import ASCENDING, DESCENDING
 
-from .functions import SearchScore
+from .functions import SearchExpression
 from .query import MongoQuery, wrap_database_errors
 
 
@@ -115,31 +115,31 @@ def _prepare_search_expressions_for_pipeline(self, expression, target, search_id
         for sub_expr in self._get_search_expressions(expression):
             alias = f"__search_expr.search{next(search_idx)}"
             replacements[sub_expr] = self._get_replace_expr(sub_expr, searches, alias)
-        return replacements, searches
+        return replacements, list(searches.values())
 
     def _prepare_search_query_for_aggregation_pipeline(self, order_by):
         replacements = {}
-        searches = {}
-        search_idx = itertools.count(start=1)
+        searches = []
+        annotation_group_idx = itertools.count(start=1)
         for target, expr in self.query.annotation_select.items():
             new_replacements, expr_searches = self._prepare_search_expressions_for_pipeline(
-                expr, target, search_idx
+                expr, target, annotation_group_idx
             )
             replacements.update(new_replacements)
-            searches.update(expr_searches)
+            searches += expr_searches
 
         for expr, _ in order_by:
             new_replacements, expr_searches = self._prepare_search_expressions_for_pipeline(
-                expr, None, search_idx
+                expr, None, annotation_group_idx
             )
             replacements.update(new_replacements)
-            searches.update(expr_searches)
+            searches += expr_searches
 
         having_replacements, having_group = self._prepare_search_expressions_for_pipeline(
-            self.having, None, search_idx
+            self.having, None, annotation_group_idx
         )
         replacements.update(having_replacements)
-        searches.update(having_group)
+        searches += having_group
         return searches, replacements
 
     def _prepare_annotations_for_aggregation_pipeline(self, order_by):
@@ -249,14 +249,21 @@ def _build_aggregation_pipeline(self, ids, group):
                 pipeline.append({"$unset": "_id"})
         return pipeline
 
+    def _compound_searches_queries(self, searches):
+        if not searches:
+            return []
+        if len(searches) > 1:
+            raise ValueError("Cannot perform more than one search operation.")
+        return [searches[0], {"$addFields": {"__search_expr.search1": {"$meta": "searchScore"}}}]
+
     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
         )
         group, group_replacements = self._prepare_annotations_for_aggregation_pipeline(order_by)
         all_replacements = {**search_replacements, **group_replacements}
-        self.search_pipeline = searches
+        self.search_pipeline = self._compound_searches_queries(searches)
         # query.group_by is either:
         # - None: no GROUP BY
         # - True: group by select fields
@@ -607,7 +614,7 @@ def _get_aggregate_expressions(self, expr):
         return self._get_all_expressions_of_type(expr, Aggregate)
 
     def _get_search_expressions(self, expr):
-        return self._get_all_expressions_of_type(expr, SearchScore)
+        return self._get_all_expressions_of_type(expr, SearchExpression)
 
     def _get_all_expressions_of_type(self, expr, target_type):
         stack = [expr]

django_mongodb_backend/functions.py

@@ -3,7 +3,7 @@
 from django.conf import settings
 from django.db import NotSupportedError
 from django.db.models import DateField, DateTimeField, Expression, FloatField, TimeField
-from django.db.models.expressions import F, Func, Value
+from django.db.models.expressions import Func
 from django.db.models.functions import JSONArray
 from django.db.models.functions.comparison import Cast, Coalesce, Greatest, Least, NullIf
 from django.db.models.functions.datetime import (
@@ -38,9 +38,8 @@
     Trim,
     Upper,
 )
-from django.utils.deconstruct import deconstructible
 
-from .query_utils import process_lhs, process_rhs
+from .query_utils import process_lhs
 
 MONGO_OPERATORS = {
     Ceil: "ceil",
@@ -269,28 +268,160 @@ def trunc_time(self, compiler, connection):
     }
 
 
-@deconstructible(path="django_mongodb_backend.functions.SearchScore")
-class SearchScore(Expression):
-    def __init__(self, path, value, operation="equals", **kwargs):
-        self.extra_params = kwargs
-        self.lhs = path if hasattr(path, "resolve_expression") else F(path)
-        if not isinstance(value, str):
-            # TODO HANDLE VALUES LIKE Value("some string")
-            raise ValueError("STRING NEEDED")
-        self.rhs = Value(value)
-        self.operation = operation
+class SearchExpression(Expression):
+    optional_arguments = []
+
+    def __init__(self, *args, score=None, **kwargs):
+        self.score = score
+        # Support positional arguments first
+        if args and len(args) > len(self.expected_arguments) + len(self.optional_arguments):
+            raise ValueError(
+                f"Too many positional arguments: expected {len(self.expected_arguments)}"
+            )
+        # TODO: REFACTOR.
+        for arg_name, arg_type in self.expected_arguments:
+            if args:
+                value = args.pop(0)
+                if arg_name in kwargs:
+                    raise ValueError(
+                        f"Argument '{arg_name}' was provided both positionally and as keyword"
+                    )
+            elif arg_name in kwargs:
+                value = kwargs.pop(arg_name)
+            else:
+                raise ValueError(f"Missing required argument '{arg_name}'")
+            if not isinstance(value, arg_type):
+                raise ValueError(f"Argument '{arg_name}' must be of type {arg_type.__name__}")
+            setattr(self, arg_name, value)
+
+        for arg_name, arg_type in self.optional_arguments:
+            if args:
+                if arg_name in kwargs:
+                    raise ValueError(
+                        f"Argument '{arg_name}' was provided both positionally and as keyword"
+                    )
+                value = args.pop(0)
+            elif arg_name in kwargs:
+                value = kwargs.pop(arg_name)
+            else:
+                value = None
+            if value is not None and not isinstance(value, arg_type):
+                raise ValueError(f"Argument '{arg_name}' must be of type {arg_type.__name__}")
+            setattr(self, arg_name, value)
+        if kwargs:
+            raise ValueError(f"Unexpected keyword arguments: {list(kwargs.keys())}")
         super().__init__(output_field=FloatField())
 
+    def get_source_expressions(self):
+        return []
+
+    def __str__(self):
+        args = ", ".join(map(str, self.get_source_expressions()))
+        return f"{self.search_type}({args})"
+
     def __repr__(self):
-        return f"search {self.field} = {self.value} | {self.extra_params}"
+        return str(self)
+
+    def as_sql(self, compiler, connection):
+        return "", []
+
+    def _get_query_index(self, field, compiler):
+        for search_indexes in compiler.collection.list_search_indexes():
+            mappings = search_indexes["latestDefinition"]["mappings"]
+            if mappings["dynamic"] or field in mappings["fields"]:
+                return search_indexes["name"]
+        return "default"
 
     def as_mql(self, compiler, connection):
-        lhs = process_lhs(self, compiler, connection)
-        rhs = process_rhs(self, compiler, connection)
-        return {"$search": {self.operation: {"path": lhs[:1], "query": rhs, **self.extra_params}}}
+        params = {}
+        for arg_name, _ in self.expected_arguments:
+            params[arg_name] = getattr(self, arg_name)
+        if self.score:
+            params["score"] = self.score.as_mql(compiler, connection)
+        index = self._get_query_index(params.get("path"), compiler)
+        return {"$search": {self.search_type: params, "index": index}}
+
+
+class SearchAutocomplete(SearchExpression):
+    search_type = "autocomplete"
+    expected_arguments = [("path", str), ("query", str)]
+
+
+class SearchEquals(SearchExpression):
+    search_type = "equals"
+    expected_arguments = [("path", str), ("value", str)]
+
+
+class SearchExists(SearchExpression):
+    search_type = "equals"
+    expected_arguments = [("path", str)]
+
+
+class SearchIn(SearchExpression):
+    search_type = "equals"
+    expected_arguments = [("path", str), ("value", str | list)]
+
+
+class SearchPhrase(SearchExpression):
+    search_type = "equals"
+    expected_arguments = [("path", str), ("value", str | list)]
+    optional_arguments = [("slop", int), ("synonyms", str)]
+
+
+"""
+IT IS BEING REFACTORED
+class SearchOperator(SearchExpression):
+    _operation_params = {
+        "autocomplete": ("path", {"query"}),
+        "equals": ("path", {"value"}),
+        "exists": ("path", {}),
+        "in": ("path", {"value"}),
+        "phrase": ("path", {"query"}),
+        "queryString": ("defaultPath", {"query"}),
+        "range": ("path", {("lt", "lte"), ("gt", "gte")}),
+        "regex": ("path", {"query"}),
+        "text": ("path", {"query"}),
+        "wildcard": ("path", {"query"}),
+        "geoShape": ("path", {"query", "relation", "geometry"}),
+        "geoWithin": ("path", {("box", "circle", "geometry")}),
+        "moreLikeThis": (None, {"like"}),
+        "near": ("path", {"origin", "pivot"}),
+    }
+
+    def __init__(self, operation, **kwargs):
+        self.lhs = path if path is None or hasattr(path, "resolve_expression") else F(path)
+        self.operation = operation
+        self.lhs_field, needed_params = self._operation_params[self.operation]
+        rhs_values = {}
+        for param in needed_params:
+            if isinstance(param, str):
+                rhs_values[param] = kwargs.pop(param)
+            else:
+                for key in param:
+                    if key in kwargs:
+                        rhs_values[param] = kwargs.pop(key)
+                        break
+                else:
+                    raise ValueError(f"Not found either {', '.join(param)}")
+
+        self.rhs_values = rhs_values
+        self.extra_params = kwargs
+        super().__init__(output_field=FloatField())
+
+    def as_mql(self, compiler, connection):
+        params = {**self.rhs_values, **self.extra_params}
+        if self.lhs:
+            lhs_mql = process_lhs(self, compiler, connection)
+            params[self.lhs_field] = lhs_mql[1:]
+        index = self._get_query_index(compiler, connection)
+        return {"$search": {self.operation: params, "index": index}}
+
+    def get_source_expressions(self):
+        return [self.lhs, self.rhs, self.extra_params]
 
     def as_sql(self, compiler, connection):
         return "", []
+"""
 
 
 def register_functions():

tests/queries_/test_search.py

@@ -0,0 +1,12 @@
+from django.test import TestCase
+
+from .models import Article
+
+from django_mongodb_backend.functions import SearchEquals
+
+
+class SearchTests(TestCase):
+    def test_1(self):
+        Article.objects.create(headline="cross", number=1, body="body")
+        aa = Article.objects.annotate(score=SearchEquals(path="headline", value="cross")).all()
+        print(aa)