[WIP] Atlas search lookups

.pre-commit-config.yaml

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

django_mongodb_backend/__init__.py

@@ -8,7 +8,7 @@
 
 from .aggregates import register_aggregates  # noqa: E402
 from .checks import register_checks  # noqa: E402
-from .expressions import register_expressions  # noqa: E402
+from .expressions.builtins import register_expressions  # noqa: E402
 from .fields import register_fields  # noqa: E402
 from .functions import register_functions  # noqa: E402
 from .indexes import register_indexes  # noqa: E402

django_mongodb_backend/compiler.py

@@ -17,6 +17,7 @@
 from django.utils.functional import cached_property
 from pymongo import ASCENDING, DESCENDING
 
+from .expressions.search import SearchExpression, SearchVector
 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,30 @@ 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, 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)
+
+    def _prepare_search_query_for_aggregation_pipeline(self, order_by):
+        replacements = {}
+        annotation_group_idx = itertools.count(start=1)
+        for expr in self.query.annotation_select.values():
+            self._prepare_search_expressions_for_pipeline(expr, annotation_group_idx, replacements)
+
+        for expr, _ in order_by:
+            self._prepare_search_expressions_for_pipeline(expr, annotation_group_idx, replacements)
+
+        self._prepare_search_expressions_for_pipeline(
+            self.having, annotation_group_idx, replacements
+        )
+        return replacements
+
     def _prepare_annotations_for_aggregation_pipeline(self, order_by):
         """Prepare annotations for the aggregation pipeline."""
         replacements = {}
@@ -179,6 +206,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 = []
@@ -207,9 +237,36 @@ def _build_aggregation_pipeline(self, ids, group):
                 pipeline.append({"$unset": "_id"})
         return pipeline
 
+    def _compound_searches_queries(self, search_replacements):
+        if not search_replacements:
+            return []
+        if len(search_replacements) > 1:
+            raise ValueError("Cannot perform more than one search operation.")
+        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)
-        group, all_replacements = self._prepare_annotations_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(search_replacements)
         # query.group_by is either:
         # - None: no GROUP BY
         # - True: group by select fields
@@ -557,10 +614,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, SearchExpression)
+
+    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/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):

django_mongodb_backend/expressions.py → django_mongodb_backend/expressions/builtins.py

@@ -26,7 +26,7 @@
 )
 from django.db.models.sql import Query
 
-from .query_utils import process_lhs
+from ..query_utils import process_lhs
 
 
 def case(self, compiler, connection):