First approach.

Author: WaVEV

django_mongodb_backend/compiler.py

@@ -17,6 +17,7 @@
 from django.utils.functional import cached_property
 from pymongo import ASCENDING, DESCENDING
 
+from .functions import SearchScore
 from .query import MongoQuery, wrap_database_errors
 
 
@@ -34,6 +35,8 @@ def __init__(self, *args, **kwargs):
         # A list of OrderBy objects for this query.
         self.order_by_objs = None
         self.subqueries = []
+        # Atlas search calls
+        self.search_pipeline = []
 
     def _get_group_alias_column(self, expr, annotation_group_idx):
         """Generate a dummy field for use in the ids fields in $group."""
@@ -57,6 +60,29 @@ def _get_column_from_expression(self, expr, alias):
         column_target.set_attributes_from_name(alias)
         return Col(self.collection_name, column_target)
 
+    def _get_replace_expr(self, sub_expr, group, alias):
+        column_target = sub_expr.output_field.clone()
+        column_target.db_column = alias
+        column_target.set_attributes_from_name(alias)
+        inner_column = Col(self.collection_name, column_target)
+        if getattr(sub_expr, "distinct", False):
+            # If the expression should return distinct values, use
+            # $addToSet to deduplicate.
+            rhs = sub_expr.as_mql(self, self.connection, resolve_inner_expression=True)
+            group[alias] = {"$addToSet": rhs}
+            replacing_expr = sub_expr.copy()
+            replacing_expr.set_source_expressions([inner_column, None])
+        else:
+            group[alias] = sub_expr.as_mql(self, self.connection)
+            replacing_expr = inner_column
+        # Count must return 0 rather than null.
+        if isinstance(sub_expr, Count):
+            replacing_expr = Coalesce(replacing_expr, 0)
+        # Variance = StdDev^2
+        if isinstance(sub_expr, Variance):
+            replacing_expr = Power(replacing_expr, 2)
+        return replacing_expr
+
     def _prepare_expressions_for_pipeline(self, expression, target, annotation_group_idx):
         """
         Prepare expressions for the aggregation pipeline.
@@ -80,29 +106,42 @@ def _prepare_expressions_for_pipeline(self, expression, target, annotation_group
             alias = (
                 f"__aggregation{next(annotation_group_idx)}" if sub_expr != expression else target
             )
-            column_target = sub_expr.output_field.clone()
-            column_target.db_column = alias
-            column_target.set_attributes_from_name(alias)
-            inner_column = Col(self.collection_name, column_target)
-            if sub_expr.distinct:
-                # If the expression should return distinct values, use
-                # $addToSet to deduplicate.
-                rhs = sub_expr.as_mql(self, self.connection, resolve_inner_expression=True)
-                group[alias] = {"$addToSet": rhs}
-                replacing_expr = sub_expr.copy()
-                replacing_expr.set_source_expressions([inner_column, None])
-            else:
-                group[alias] = sub_expr.as_mql(self, self.connection)
-                replacing_expr = inner_column
-            # Count must return 0 rather than null.
-            if isinstance(sub_expr, Count):
-                replacing_expr = Coalesce(replacing_expr, 0)
-            # Variance = StdDev^2
-            if isinstance(sub_expr, Variance):
-                replacing_expr = Power(replacing_expr, 2)
-            replacements[sub_expr] = replacing_expr
+            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):
+        searches = {}
+        replacements = {}
+        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
+
+    def _prepare_search_query_for_aggregation_pipeline(self, order_by):
+        replacements = {}
+        searches = {}
+        search_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
+            )
+            replacements.update(new_replacements)
+            searches.update(expr_searches)
+
+        for expr, _ in order_by:
+            new_replacements, expr_searches = self._prepare_search_expressions_for_pipeline(
+                expr, None, search_idx
+            )
+            replacements.update(new_replacements)
+            searches.update(expr_searches)
+
+        having_replacements, having_group = self._prepare_search_expressions_for_pipeline(
+            self.having, None, search_idx
+        )
+        replacements.update(having_replacements)
+        searches.update(having_group)
+        return searches, replacements
+
     def _prepare_annotations_for_aggregation_pipeline(self, order_by):
         """Prepare annotations for the aggregation pipeline."""
         replacements = {}
@@ -179,6 +218,9 @@ def _get_group_id_expressions(self, order_by):
             ids = self.get_project_fields(tuple(columns), force_expression=True)
         return ids, replacements
 
+    def _build_search_pipeline(self, search_queries):
+        pass
+
     def _build_aggregation_pipeline(self, ids, group):
         """Build the aggregation pipeline for grouping."""
         pipeline = []
@@ -209,7 +251,12 @@ def _build_aggregation_pipeline(self, ids, group):
 
     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)
-        group, all_replacements = self._prepare_annotations_for_aggregation_pipeline(order_by)
+        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
         # query.group_by is either:
         # - None: no GROUP BY
         # - True: group by select fields
@@ -557,10 +604,16 @@ def get_lookup_pipeline(self):
         return result
 
     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)
+
+    def _get_all_expressions_of_type(self, expr, target_type):
         stack = [expr]
         while stack:
             expr = stack.pop()
-            if isinstance(expr, Aggregate):
+            if isinstance(expr, target_type):
                 yield expr
             elif hasattr(expr, "get_source_expressions"):
                 stack.extend(expr.get_source_expressions())

django_mongodb_backend/functions.py

@@ -2,8 +2,8 @@
 
 from django.conf import settings
 from django.db import NotSupportedError
-from django.db.models import DateField, DateTimeField, TimeField
-from django.db.models.expressions import Func
+from django.db.models import DateField, DateTimeField, Expression, FloatField, TimeField
+from django.db.models.expressions import F, Func, Value
 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,8 +38,9 @@
     Trim,
     Upper,
 )
+from django.utils.deconstruct import deconstructible
 
-from .query_utils import process_lhs
+from .query_utils import process_lhs, process_rhs
 
 MONGO_OPERATORS = {
     Ceil: "ceil",
@@ -268,6 +269,30 @@ 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
+        super().__init__(output_field=FloatField())
+
+    def __repr__(self):
+        return f"search {self.field} = {self.value} | {self.extra_params}"
+
+    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}}}
+
+    def as_sql(self, compiler, connection):
+        return "", []
+
+
 def register_functions():
     Cast.as_mql = cast
     Concat.as_mql = concat

django_mongodb_backend/query.py

@@ -49,6 +49,7 @@ def __init__(self, compiler):
         self.lookup_pipeline = None
         self.project_fields = None
         self.aggregation_pipeline = compiler.aggregation_pipeline
+        self.search_pipeline = compiler.search_pipeline
         self.extra_fields = None
         self.combinator_pipeline = None
         # $lookup stage that encapsulates the pipeline for performing a nested
@@ -75,6 +76,8 @@ def get_cursor(self):
 
     def get_pipeline(self):
         pipeline = []
+        if self.search_pipeline:
+            pipeline.extend(self.search_pipeline)
         if self.lookup_pipeline:
             pipeline.extend(self.lookup_pipeline)
         for query in self.subqueries or ():

django_mongodb_backend/query_utils.py

@@ -4,7 +4,7 @@
 
 
 def is_direct_value(node):
-    return not hasattr(node, "as_sql")
+    return not hasattr(node, "as_sql") and not hasattr(node, "as_mql")
 
 
 def process_lhs(node, compiler, connection):

tests/queries_/models.py

@@ -1,6 +1,7 @@
 from django.db import models
 
 from django_mongodb_backend.fields import ObjectIdAutoField, ObjectIdField
+from django_mongodb_backend.indexes import SearchIndex
 
 
 class Author(models.Model):
@@ -53,3 +54,12 @@ class Meta:
 
     def __str__(self):
         return str(self.pk)
+
+
+class Article(models.Model):
+    headline = models.CharField(max_length=100)
+    number = models.IntegerField()
+    body = models.TextField()
+
+    class Meta:
+        indexes = [SearchIndex(fields=["headline"])]