Total doc count added to agg pipeline endpoints

materialsproject · Jun 27, 2021 · 78b1e86 · 78b1e86
1 parent 6b3762f
commit 78b1e86
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 16 deletions.
diff --git a/src/mp_api/routes/robocrys/query_operators.py b/src/mp_api/routes/robocrys/query_operators.py
@@ -38,13 +38,29 @@ def query(
                 }
             },
             {
-                "$project": {
-                    "_id": 0,
-                    "task_id": 1,
-                    "description": 1,
-                    "condensed_structure": 1,
-                    "last_updates": 1,
-                    "search_score": {"$meta": "searchScore"},
+                "$facet": {
+                    "total_doc": [{"$count": "count"}],
+                    "results": [
+                        {
+                            "$project": {
+                                "_id": 0,
+                                "task_id": 1,
+                                "description": 1,
+                                "condensed_structure": 1,
+                                "last_updates": 1,
+                                "search_score": {"$meta": "searchScore"},
+                            }
+                        }
+                    ],
+                }
+            },
+            {"$unwind": "$results"},
+            {"$unwind": "$total_doc"},
+            {
+                "$replaceRoot": {
+                    "newRoot": {
+                        "$mergeObjects": ["$results", {"total_doc": "$total_doc.count"}]
+                    }
                 }
             },
             {"$sort": {"search_score": -1}},
@@ -53,5 +69,13 @@ def query(
         ]
         return {"pipeline": pipeline}
 
+    def post_process(self, docs):
+        print(docs[0])
+        self.total_doc = docs[0]["total_doc"]
+        return docs
+
+    def meta(self):
+        return {"total_doc": self.total_doc}
+
     def ensure_indexes(self):
         return [("description", False)]
diff --git a/src/mp_api/routes/synthesis/query_operators.py b/src/mp_api/routes/synthesis/query_operators.py
@@ -98,7 +98,8 @@ def query(
                     "highlights": {"$meta": "searchHighlights"},
                 }
             )
-        pipeline.append({"$project": project_dict})
+        pipeline.append({"$facet": {"results": [], "total_doc": []}})
+        pipeline[-1]["$facet"]["results"].append({"$project": project_dict})
 
         crit: Dict[str, Any] = {}
         if synthesis_type:
@@ -145,24 +146,40 @@ def query(
             }
 
         if crit:
-            pipeline.append({"$match": crit})
-
-            # total = next(
-            #     self.store._collection.aggregate(
-            #         pipeline + [{"$count": "total"}], allowDiskUse=True
-            #     )
-            # )["total"]
+            pipeline[-1]["$facet"]["results"].append({"$match": crit})
 
+        pipeline[-1]["$facet"]["total_doc"].append({"$count": "count"})
         pipeline.extend(
-            [{"$sort": {"search_score": -1}}, {"$skip": skip}, {"$limit": limit}]
+            [
+                {"$unwind": "$results"},
+                {"$unwind": "$total_doc"},
+                {
+                    "$replaceRoot": {
+                        "newRoot": {
+                            "$mergeObjects": [
+                                "$results",
+                                {"total_doc": "$total_doc.count"},
+                            ]
+                        }
+                    }
+                },
+                {"$sort": {"search_score": -1}},
+                {"$skip": skip},
+                {"$limit": limit},
+            ]
         )
 
         return {"pipeline": pipeline}
 
     def post_process(self, docs):
 
+        self.total_doc = docs[0]["total_doc"]
+
         for doc in docs:
             mask_highlights(doc)
             mask_paragraphs(doc)
 
         return docs
+
+    def meta(self):
+        return {"total_doc": self.total_doc}