From 247a4bf8b2925b0192e5bcfbc5c271392917f51f Mon Sep 17 00:00:00 2001
From: Jason Munro <jason.munro@gmail.com>
Date: Fri, 9 Sep 2022 17:04:01 -0700
Subject: [PATCH] Task API endpoint enhancement (#531)

* Fix trajectory frame prop generation

* Add post-process to remove task fields

* Update task hint scheme

* Disable validation on task endpoints

* Add default return value to task doc value pop
---
 .../emmet/api/routes/tasks/hint_scheme.py     | 13 ++++++++----
 .../emmet/api/routes/tasks/query_operators.py | 12 +++++++++++
 emmet-api/emmet/api/routes/tasks/resources.py | 12 +++++++----
 emmet-api/emmet/api/routes/tasks/utils.py     | 20 ++++++++++---------
 4 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/emmet-api/emmet/api/routes/tasks/hint_scheme.py b/emmet-api/emmet/api/routes/tasks/hint_scheme.py
index ee3a4e41d7..17535a8a0b 100644
--- a/emmet-api/emmet/api/routes/tasks/hint_scheme.py
+++ b/emmet-api/emmet/api/routes/tasks/hint_scheme.py
@@ -8,9 +8,14 @@ class TasksHintScheme(HintScheme):
 
     def generate_hints(self, query):
 
-        hints = {"hint": {}}
-
         if query["criteria"] == {}:
-            hints["hint"]["_id"] = 1
+            return {"hint": {"_id": 1}}
+
+        for param in query["criteria"]:
+
+            if "composition_reduced" in param:
+                return {"hint": {"composition_reduced.$**": 1}}
+            elif "nelements" in param:
+                return {"hint": {"nelements": 1}}
 
-        return hints
+        return {"hint": {}}
diff --git a/emmet-api/emmet/api/routes/tasks/query_operators.py b/emmet-api/emmet/api/routes/tasks/query_operators.py
index 7153e1c53e..0065df326b 100644
--- a/emmet-api/emmet/api/routes/tasks/query_operators.py
+++ b/emmet-api/emmet/api/routes/tasks/query_operators.py
@@ -31,6 +31,18 @@ def query(
 
         return {"criteria": crit}
 
+    def post_process(self, docs, query):
+        """
+        Post processing to remove unwanted fields from all task queries
+        """
+
+        for doc in docs:
+            doc.pop("tags", None)
+            doc.pop("sbxn", None)
+            doc.pop("dir_name", None)
+
+        return docs
+
 
 class TrajectoryQuery(QueryOperator):
     """
diff --git a/emmet-api/emmet/api/routes/tasks/resources.py b/emmet-api/emmet/api/routes/tasks/resources.py
index e09fe0b81b..570eddcd4f 100644
--- a/emmet-api/emmet/api/routes/tasks/resources.py
+++ b/emmet-api/emmet/api/routes/tasks/resources.py
@@ -1,3 +1,4 @@
+from this import d
 from maggma.api.query_operator import PaginationQuery, SortQuery, SparseFieldsQuery
 from maggma.api.resource import ReadOnlyResource
 
@@ -38,7 +39,8 @@ def task_resource(task_store):
         header_processor=GlobalHeaderProcessor(),
         hint_scheme=TasksHintScheme(),
         tags=["Tasks"],
-        timeout=timeout
+        timeout=timeout,
+        disable_validation=True,
     )
 
     return resource
@@ -54,7 +56,7 @@ def task_deprecation_resource(materials_store):
         enable_default_search=True,
         sub_path="/deprecation/",
         header_processor=GlobalHeaderProcessor(),
-        timeout=timeout
+        timeout=timeout,
     )
 
     return resource
@@ -69,7 +71,8 @@ def trajectory_resource(task_store):
         tags=["Tasks"],
         sub_path="/trajectory/",
         header_processor=GlobalHeaderProcessor(),
-        timeout=timeout
+        timeout=timeout,
+        disable_validation=True,
     )
 
     return resource
@@ -92,7 +95,8 @@ def entries_resource(task_store):
         tags=["Tasks"],
         sub_path="/entries/",
         header_processor=GlobalHeaderProcessor(),
-        timeout=timeout
+        timeout=timeout,
+        disable_validation=True,
     )
 
     return resource
diff --git a/emmet-api/emmet/api/routes/tasks/utils.py b/emmet-api/emmet/api/routes/tasks/utils.py
index fa968b9d04..11a5cda7ec 100644
--- a/emmet-api/emmet/api/routes/tasks/utils.py
+++ b/emmet-api/emmet/api/routes/tasks/utils.py
@@ -22,7 +22,7 @@ def calcs_reversed_to_trajectory(calcs_reversed: List[dict]):
 
     for calculation in calcs_reversed:
         structures = []
-        frame_props = defaultdict(list)  # type: dict
+        frame_props = []
 
         steps = calculation.get("output", {}).get("ionic_steps", None)
 
@@ -33,6 +33,8 @@ def calcs_reversed_to_trajectory(calcs_reversed: List[dict]):
         else:
             for step in steps:
 
+                step_dict = {}
+
                 structure_dict = step.get("structure", None)
 
                 if structure_dict is not None:
@@ -40,14 +42,14 @@ def calcs_reversed_to_trajectory(calcs_reversed: List[dict]):
 
                 structures.append(structure)
 
-                frame_props["e_fr_energy"].append(step.get("e_fr_energy", None))
-                frame_props["e_wo_entrp"].append(step.get("e_wo_entrp", None))
-                frame_props["e_0_energy"].append(step.get("e_0_energy", None))
-                frame_props["forces"].append(step.get("forces", None))
-                frame_props["stresses"].append(step.get("stress", None))
-                frame_props["electronic_steps"].append(
-                    step.get("electronic_steps", None)
-                )
+                step_dict["e_fr_energy"] = step.get("e_fr_energy", None)
+                step_dict["e_wo_entrp"] = step.get("e_wo_entrp", None)
+                step_dict["e_0_energy"] = step.get("e_0_energy", None)
+                step_dict["forces"] = step.get("forces", None)
+                step_dict["stresses"] = step.get("stress", None)
+                step_dict["electronic_steps"] = step.get("electronic_steps", None)
+
+                frame_props.append(step_dict)
 
             traj = Trajectory.from_structures(
                 structures, frame_properties=frame_props, time_step=None