From 54cbd4dac56b1443c5ed0748c7584ea59400696a Mon Sep 17 00:00:00 2001
From: Patrick Huck <phuck@lbl.gov>
Date: Wed, 18 Dec 2024 17:46:22 -0800
Subject: [PATCH] BatchIdQuery

---
 .../materials/materials/query_operators.py    | 32 +++++++++++++++++++
 .../routes/materials/materials/resources.py   |  2 ++
 .../api/routes/materials/summary/resources.py |  2 ++
 3 files changed, 36 insertions(+)

diff --git a/emmet-api/emmet/api/routes/materials/materials/query_operators.py b/emmet-api/emmet/api/routes/materials/materials/query_operators.py
index 39eec95e18..a2b7f69095 100644
--- 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
index 917d73af39..772d0d040f 100644
--- 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
index 633617580e..d3bbd3a605 100644
--- 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(),