Refactor.

Author: WaVEV

.pre-commit-config.yaml

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

django_mongodb_backend/expressions/builtins.py

@@ -5,6 +5,7 @@
 from bson import Decimal128
 from django.core.exceptions import EmptyResultSet, FullResultSet
 from django.db import NotSupportedError
+from django.db.models import Expression, FloatField
 from django.db.models.expressions import (
     Case,
     Col,
@@ -207,6 +208,293 @@ def value(self, compiler, connection):  # noqa: ARG001
     return value
 
 
+class SearchExpression(Expression):
+    def __init__(self):
+        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 str(self)
+
+    def as_sql(self, compiler, connection):
+        return "", []
+
+    def _get_query_index(self, fields, compiler):
+        fields = set(fields)
+        for search_indexes in compiler.collection.list_search_indexes():
+            mappings = search_indexes["latestDefinition"]["mappings"]
+            if mappings["dynamic"] or fields.issubset(set(mappings["fields"])):
+                return search_indexes["name"]
+        return "default"
+
+
+class SearchAutocomplete(SearchExpression):
+    def __init__(self, path, query, score=None):
+        self.path = F(path)
+        self.query = Value(query)
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+            "query": self.query.as_mql(compiler, connection),
+        }
+        if self.score is not None:
+            params["score"] = self.score
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"autocomplete": params, "index": index}}
+
+
+class SearchEquals(SearchExpression):
+    def __init__(self, path, value, score=None):
+        self.path = F(path)
+        self.value = Value(query)
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+            "value": self.value.as_mql(compiler, connection),
+        }
+        if self.score is not None:
+            params["score"] = self.score
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"equals": params, "index": index}}
+
+
+class SearchExists(SearchExpression):
+    def __init__(self, path, score=None):
+        self.path = F(path)
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+        }
+        if self.score is not None:
+            params["score"] = self.score
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"exists": params, "index": index}}
+
+
+class SearchIn(SearchExpression):
+    def __init__(self, path, value, score=None):
+        self.path = F(path)
+        self.value = Value(value)
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+            "value": self.value.as_mql(compiler, connection),
+        }
+        if self.score is not None:
+            params["score"] = self.score
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"in": params, "index": index}}
+
+
+class SearchPhrase(SearchExpression):
+    def __init__(self, path, value, slop=None, synonyms=None, score=None):
+        self.path = F(path)
+        self.value = Value(value)
+        self.score = score
+        self.slop = slop
+        self.synonyms = synonyms
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+            "value": self.value.as_mql(compiler, connection),
+        }
+        if self.score is not None:
+            params["score"] = self.score
+        if self.slop is not None:
+            params["slop"] = self.slop
+        if self.synonyms is not None:
+            params["synonyms"] = self.synonyms
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"phrase": params, "index": index}}
+
+
+class SearchQueryString(SearchExpression):
+    def __init__(self, path, query, score=None):
+        self.path = F(path)
+        self.query = Value(query)
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "defaultPath": self.path.as_mql(compiler, connection)[1:],
+            "query": self.query.as_mql(compiler, connection),
+        }
+        if self.score is not None:
+            params["score"] = self.score
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"queryString": params, "index": index}}
+
+
+class SearchRange(SearchExpression):
+    def __init__(self, path, lt=None, lte=None, gt=None, gte=None, score=None):
+        self.path = F(path)
+        self.lt = Value(lt)
+        self.lte = Value(lte)
+        self.gt = Value(gt)
+        self.gte = Value(gte)
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+        }
+        if self.score is not None:
+            params["score"] = self.score
+        if self.lt is not None:
+            params["lt"] = self.lt.as_mql(compiler, connection)
+        if self.lte is not None:
+            params["lte"] = self.lte.as_mql(compiler, connection)
+        if self.gt is not None:
+            params["gt"] = self.gt.as_mql(compiler, connection)
+        if self.gte is not None:
+            params["gte"] = self.gte.as_mql(compiler, connection)
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"range": params, "index": index}}
+
+
+class SearchRegex(SearchExpression):
+    def __init__(self, path, query, allow_analyzed_field=None, score=None):
+        self.path = F(path)
+        self.allow_analyzed_field = Value(allow_analyzed_field)
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+        }
+        if self.score:
+            params["score"] = self.score
+        if self.allow_analyzed_field is not None:
+            params["allowAnalyzedField"] = self.allow_analyzed_field.as_mql(compiler, connection)
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"regex": params, "index": index}}
+
+
+class SearchText(SearchExpression):
+    def __init__(self, path, query, fuzzy=None, match_criteria=None, synonyms=None, score=None):
+        self.path = F(path)
+        self.fuzzy = Value(fuzzy)
+        self.match_criteria = Value(match_criteria)
+        self.synonyms = Value(synonyms)
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+        }
+        if self.score:
+            params["score"] = self.score
+        if self.fuzzy is not None:
+            params["fuzzy"] = self.fuzzy.as_mql(compiler, connection)
+        if self.match_criteria is not None:
+            params["matchCriteria"] = self.match_criteria.as_mql(compiler, connection)
+        if self.synonyms is not None:
+            params["synonyms"] = self.synonyms.as_mql(compiler, connection)
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"text": params, "index": index}}
+
+
+class SearchWildcard(SearchExpression):
+    def __init__(self, path, query, allow_analyzed_field=None, score=None):
+        self.path = F(path)
+        self.allow_analyzed_field = Value(allow_analyzed_field)
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+        }
+        if self.score:
+            params["score"] = self.score
+        if self.allow_analyzed_field is not None:
+            params["allowAnalyzedField"] = self.allow_analyzed_field.as_mql(compiler, connection)
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"wildcard": params, "index": index}}
+
+
+class SearchGeoShape(SearchExpression):
+    def __init__(self, path, relation, geometry, score=None):
+        self.path = F(path)
+        self.relation = relation
+        self.geometry = geometry
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+            "relation": self.relation,
+            "geometry": self.geometry,
+        }
+        if self.score:
+            params["score"] = self.score
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"wildcard": params, "index": index}}
+
+
+class SearchGeoWithin(SearchExpression):
+    def __init__(self, path, kind, geo_object, geometry, score=None):
+        self.path = F(path)
+        self.kind = kind
+        self.geo_object = geo_object
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection)[1:],
+            self.kind: self.geo_object,
+        }
+        if self.score:
+            params["score"] = self.score
+        index = self._get_query_index([self.path], compiler)
+        return {"$search": {"wildcard": params, "index": index}}
+
+
+class SearchMoreLikeThis(SearchExpression):
+    def __init__(self, documents, score=None):
+        self.documents = documents
+        self.score = score
+        super().__init__()
+
+    def as_mql(self, compiler, connection):
+        params = {
+            "like": self.documents,
+        }
+        if self.score:
+            params["score"] = self.score
+        needed_fields = []
+        for doc in self.documents:
+            needed_fields += list(doc.keys())
+        index = self._get_query_index(needed_fields, compiler)
+        return {"$search": {"wildcard": params, "index": index}}
+
+
 def register_expressions():
     Case.as_mql = case
     Col.as_mql = col