Refactor

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
 
 
@@ -117,29 +117,24 @@ def _prepare_search_expressions_for_pipeline(
             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(
+            self._prepare_search_expressions_for_pipeline(
                 expr, target, annotation_group_idx, replacements
             )
-            searches += expr_searches
 
         for expr, _ in order_by:
-            expr_searches = self._prepare_search_expressions_for_pipeline(
+            self._prepare_search_expressions_for_pipeline(
                 expr, None, annotation_group_idx, replacements
             )
-            searches += expr_searches
 
-        having_group = self._prepare_search_expressions_for_pipeline(
+        self._prepare_search_expressions_for_pipeline(
             self.having, None, 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 +243,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