Configure query on request (#1073)

* changes for docker on arm64 architecture * add header - query translation * add none to allowed license * Revert "changes for docker on arm64 architecture" This reverts commit b1c674f. * slightly shorten * update * remove unnecessary comment * fix import statement * update header check and license type * disable query config function for now * lint * rm comment, simplify criteria logic * split groups, use set ops * configure license query on request * avoid mypy complain --------- Co-authored-by: Patrick Huck <[email protected]>
materialsproject · Dec 19, 2024 · 3837454 · 3837454
1 parent 7b8181e
commit 3837454
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 8 deletions.
diff --git a/emmet-api/emmet/api/core/global_header.py b/emmet-api/emmet/api/core/global_header.py
@@ -1,6 +1,7 @@
 from maggma.api.resource.core import HeaderProcessor
 from fastapi import Response, Request
-from maggma.api.query_operator import QueryOperator
+from maggma.api.utils import STORE_PARAMS
+from emmet.api.routes.materials.materials.query_operators import LicenseQuery
 
 
 class GlobalHeaderProcessor(HeaderProcessor):
@@ -16,6 +17,14 @@ def process_header(self, response: Response, request: Request):
         response.headers["X-Consumer-Id"] = consumer_id
 
     def configure_query_on_request(
-        self, request: Request, query_operator: QueryOperator
-    ):
-        pass
+        self, request: Request, query_operator: LicenseQuery
+    ) -> STORE_PARAMS:
+        groups = request.headers.get("x-consumer-groups", "")
+        if not groups:
+            return query_operator.query(license="BY-C")
+
+        grps = set(group.strip() for group in groups.split(","))
+        if grps & {"TERMS:ACCEPT-NC", "admin"}:
+            return query_operator.query(license="All")
+
+        return query_operator.query(license="BY-C")
diff --git a/emmet-api/emmet/api/routes/materials/materials/query_operators.py b/emmet-api/emmet/api/routes/materials/materials/query_operators.py
@@ -458,12 +458,13 @@ class LicenseQuery(QueryOperator):
 
     def query(
         self,
-        license: Optional[Literal["BY-C", "BY-NC"]] = Query(
+        license: Optional[Literal["BY-C", "BY-NC", "All"]] = Query(
             "BY-C",
-            description="Query by license. Either commercial or non-commercial CC-BY",
+            description="Query by license. Can be commercial or non-commercial, or both",
         ),
     ) -> STORE_PARAMS:
-        return {"criteria": {"builder_meta.license": license}}
+        q = {"$in": ["BY-C", "BY-NC"]} if license == "All" else license
+        return {"criteria": {"builder_meta.license": q}}
 
 
 class BatchIdQuery(QueryOperator):

diff --git a/emmet-api/emmet/api/routes/materials/materials/resources.py b/emmet-api/emmet/api/routes/materials/materials/resources.py
@@ -120,10 +120,10 @@ def materials_resource(materials_store):
             BatchIdQuery(),
         ],
         header_processor=GlobalHeaderProcessor(),
+        query_to_configure_on_request=LicenseQuery(),
         tags=["Materials"],
         sub_path="/core/",
         disable_validation=True,
         timeout=MAPISettings().TIMEOUT,  # type: ignore
     )
-
     return resource
diff --git a/emmet-api/emmet/api/routes/materials/summary/resources.py b/emmet-api/emmet/api/routes/materials/summary/resources.py
@@ -73,6 +73,7 @@ def summary_resource(summary_store):
         ],
         hint_scheme=SummaryHintScheme(),
         header_processor=GlobalHeaderProcessor(),
+        query_to_configure_on_request=LicenseQuery(),
         tags=["Materials Summary"],
         sub_path="/summary/",
         disable_validation=True,