BatchIdQuery

materialsproject · Dec 19, 2024 · 54cbd4d · 54cbd4d
1 parent 237d2b7
commit 54cbd4d
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 0 deletions.
diff --git a/emmet-api/emmet/api/routes/materials/materials/query_operators.py b/emmet-api/emmet/api/routes/materials/materials/query_operators.py
@@ -464,3 +464,35 @@ def query(
         ),
     ) -> STORE_PARAMS:
         return {"criteria": {"builder_meta.license": license}}
+
+
+class BatchIdQuery(QueryOperator):
+    """Method to generate a query on batch_id"""
+
+    def query(
+        self,
+        batch_id: Optional[str] = Query(
+            None,
+            description="Query by batch identifier",
+        ),
+        batch_id_not_eq: Optional[str] = Query(
+            None,
+            description="Exclude batch identifier",
+        ),
+    ) -> STORE_PARAMS:
+        if batch_id and batch_id_not_eq:
+            raise HTTPException(
+                status_code=400,
+                detail="Please only choose one of `batch_id` and `batch_id_not_eq`.",
+            )
+
+        crit = {}  # type: dict
+        if batch_id:
+            crit["builder_meta.batch_id"] = batch_id
+        elif batch_id_not_eq:
+            crit["builder_meta.batch_id"] = {"$ne": batch_id_not_eq}
+
+        return {"criteria": crit}
+
+    def ensure_indexes(self):  # pragma: no cover
+        return [("builder_meta.batch_id", False)]
diff --git a/emmet-api/emmet/api/routes/materials/materials/resources.py b/emmet-api/emmet/api/routes/materials/materials/resources.py
@@ -25,6 +25,7 @@
     MultiMaterialIDQuery,
     LicenseQuery,
     BlessedCalcsQuery,
+    BatchIdQuery,
 )
 from emmet.api.core.global_header import GlobalHeaderProcessor
 from emmet.api.core.settings import MAPISettings
@@ -116,6 +117,7 @@ def materials_resource(materials_store):
                 default_fields=["material_id", "formula_pretty", "last_updated"],
             ),
             LicenseQuery(),
+            BatchIdQuery(),
         ],
         header_processor=GlobalHeaderProcessor(),
         tags=["Materials"],

diff --git a/emmet-api/emmet/api/routes/materials/summary/resources.py b/emmet-api/emmet/api/routes/materials/summary/resources.py
@@ -14,6 +14,7 @@
     ChemsysQuery,
     SymmetryQuery,
     LicenseQuery,
+    BatchIdQuery,
 )
 from emmet.api.routes.materials.oxidation_states.query_operators import (
     PossibleOxiStateQuery,
@@ -68,6 +69,7 @@ def summary_resource(summary_store):
             SparseFieldsQuery(SummaryDoc, default_fields=["material_id"]),
             LicenseQuery(),
             SortQuery(fields=sort_fields, max_num=1),
+            BatchIdQuery(),
         ],
         hint_scheme=SummaryHintScheme(),
         header_processor=GlobalHeaderProcessor(),