emmet-api 0.87.0.dev10__tar.gz → 0.87.0.dev11__tar.gz

{emmet_api-0.87.0.dev10 → emmet_api-0.87.0.dev11}/Dockerfile

@@ -18,7 +18,8 @@ RUN wget -q https://raw.githubusercontent.com/vishnubob/wait-for-it/master/wait-
   chmod +x wait-for-it.sh && mv wait-for-it.sh /root/.local/bin/
 
 FROM base
-ARG BUILDARCH=x86_64
+ARG BUILDARCH
+ENV BUILDARCH=x86_64
 COPY --from=builder /root/.local/lib/python3.11/site-packages /root/.local/lib/python3.11/site-packages
 COPY --from=builder /root/.local/bin /root/.local/bin
 COPY --from=builder /usr/lib/$BUILDARCH-linux-gnu/libsnappy* /usr/lib/$BUILDARCH-linux-gnu/

{emmet_api-0.87.0.dev10 → emmet_api-0.87.0.dev11}/PKG-INFO

@@ -1,17 +1,17 @@
 Metadata-Version: 2.4
 Name: emmet-api
-Version: 0.87.0.dev10
+Version: 0.87.0.dev11
 Summary: Emmet API Server
 Author-email: The Materials Project <feedback@materialsproject.org>
-License: Modified BSD
+License-Expression: BSD-3-Clause-LBNL
 Project-URL: Homepage, https://github.com/materialsproject/emmet/tree/main/emmet-api/
 Project-URL: Documentation, https://materialsproject.github.io/emmet/
 Requires-Python: >=3.11
-Requires-Dist: emmet-core>=0.86.1
+Requires-Dist: emmet-core<0.87.0,>=0.86.3rc1
 Requires-Dist: fastapi>=0.123.1
-Requires-Dist: gunicorn
+Requires-Dist: gunicorn==24.1.1
 Requires-Dist: boto3
-Requires-Dist: ddtrace
+Requires-Dist: ddtrace==4.3.0
 Requires-Dist: setproctitle
 Requires-Dist: shapely
 Requires-Dist: asgi-logger

{emmet_api-0.87.0.dev10 → emmet_api-0.87.0.dev11}/emmet/api/query_operator/__init__.py

@@ -1,6 +1,6 @@
 from emmet.api.query_operator.core import QueryOperator
 from emmet.api.query_operator.dynamic import NumericQuery, StringQueryOperator
-from emmet.api.query_operator.pagination import PaginationQuery
+from emmet.api.query_operator.pagination import AtlasPaginationQuery, PaginationQuery
 from emmet.api.query_operator.sorting import SortQuery
 from emmet.api.query_operator.sparse_fields import SparseFieldsQuery
 from emmet.api.query_operator.submission import SubmissionQuery
@@ -9,6 +9,7 @@ __all__ = [
     "QueryOperator",
     "NumericQuery",
     "StringQueryOperator",
+    "AtlasPaginationQuery",
     "PaginationQuery",
     "SortQuery",
     "SparseFieldsQuery",

{emmet_api-0.87.0.dev10 → emmet_api-0.87.0.dev11}/emmet/api/query_operator/pagination.py

@@ -4,6 +4,70 @@ from emmet.api.query_operator import QueryOperator
 from emmet.api.utils import STORE_PARAMS
 
 
+class AtlasPaginationQuery(QueryOperator):
+    """Query operators to provides pagination for Atlas Search queries."""
+
+    def __init__(self, default_limit: int = 100, max_limit: int = 1000):
+        """
+        Args:
+            default_limit: the default number of documents to return
+            max_limit: max number of documents to return.
+        """
+        self.default_limit = default_limit
+        self.max_limit = max_limit
+
+        def query(
+            _forward: bool = Query(
+                True,
+                description="Whether to page forward (True) or backward (False) in the search results.",
+            ),
+            _pagination_token: str = Query(
+                None, description="Pagination token for the next set of results."
+            ),
+            _skip: int = Query(
+                0,
+                description="Number of entries to skip in the search.",
+            ),
+            _limit: int = Query(
+                default_limit,
+                description=f"Max number of entries to return in a single query. Limited to {max_limit}.",
+            ),
+        ) -> STORE_PARAMS:
+            """
+            Pagination parameters for the API Endpoint.
+            """
+            if _limit > max_limit:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Requested more data per query than allowed by this endpoint."
+                    f" The max limit is {max_limit} entries",
+                )
+
+            if _skip < 0 or _limit < 0:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Cannot request negative _skip or _limit values",
+                )
+
+            return {
+                "forward": _forward,
+                "pagination_token": _pagination_token,
+                "skip": _skip,
+                "limit": _limit,
+            }
+
+        self.query = query  # type: ignore
+
+    def query(self):
+        """Stub query function for abstract class."""
+
+    def meta(self) -> dict:
+        """
+        Metadata for the pagination params.
+        """
+        return {"max_limit": self.max_limit}
+
+
 class PaginationQuery(QueryOperator):
     """Query operators to provides Pagination."""
 

{emmet_api-0.87.0.dev10 → emmet_api-0.87.0.dev11}/emmet/api/resource/__init__.py

@@ -12,6 +12,7 @@ from emmet.api.resource.aggregation import AggregationResource
 from emmet.api.resource.post_resource import PostOnlyResource
 from emmet.api.resource.read_resource import ReadOnlyResource, attach_query_ops
 from emmet.api.resource.submission import SubmissionResource
+from emmet.api.resource.search_resource import SearchResource
 
 __all__ = [
     "Resource",
@@ -23,4 +24,5 @@ __all__ = [
     "ReadOnlyResource",
     "attach_query_ops",
     "SubmissionResource",
+    "SearchResource",
 ]

emmet_api-0.87.0.dev11/emmet/api/resource/search_resource.py

@@ -0,0 +1,144 @@
+from inspect import signature
+from typing import Any
+
+import orjson
+from fastapi import HTTPException, Request, Response
+from pymongo.errors import NetworkTimeout, PyMongoError
+
+from emmet.api.models import Meta
+from emmet.api.query_operator import QueryOperator
+from emmet.api.resource import CollectionResource
+from emmet.api.resource.utils import (
+    attach_query_ops,
+    generate_atlas_search_pipeline,
+)
+from emmet.api.utils import STORE_PARAMS, merge_atlas_queries, serialization_helper
+
+
+class SearchResource(CollectionResource):
+    """
+    Implements a REST Compatible Resource as a GET URL endpoint
+    This class provides a number of convenience features
+    including full pagination, field projection.
+    """
+
+    def __init__(
+        self,
+        *args,
+        disable_validation: bool = False,
+        enable_default_search: bool = True,
+        query_to_configure_on_request: QueryOperator | None = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            disable_validation: Whether to use ORJSON and provide a direct FastAPI response.
+                Note this will disable auto JSON serialization and response validation with the
+                provided model.
+            enable_default_search: Enable default endpoint search behavior.
+            query_to_configure_on_request: Query operator to configure on request
+        """
+        self.disable_validation = disable_validation
+        self.enable_default_search = enable_default_search
+        self.query_to_configure_on_request = query_to_configure_on_request
+
+        super().__init__(*args, **kwargs)
+
+    def prepare_endpoint(self):
+        """
+        Internal method to prepare the endpoint by setting up default handlers
+        for routes.
+        """
+        if self.enable_default_search:
+            self.build_dynamic_model_search()
+
+    def build_dynamic_model_search(self):
+        model_name = self.model.__name__
+
+        async def search(**queries: dict[str, STORE_PARAMS]) -> dict | Response:
+
+            request: Request = queries.pop("request")  # type: ignore
+            temp_response: Response = queries.pop("temp_response")  # type: ignore
+
+            if self.query_to_configure_on_request is not None:
+                # give the key name "request", arbitrary choice, as only the value gets merged into the query
+                queries["groups"] = self.header_processor.configure_query_on_request(  # type: ignore
+                    request=request, query_operator=self.query_to_configure_on_request
+                )
+            # allowed query parameters
+            query_params = [
+                entry
+                for _, i in enumerate(self.query_operators)
+                for entry in signature(i.query).parameters
+            ]
+            # check for overlap between allowed query parameters and request query parameters
+            overlap = [key for key in request.query_params if key not in query_params]
+            if any(overlap):
+                if "limit" in overlap or "skip" in overlap:
+                    raise HTTPException(
+                        status_code=400,
+                        detail="'limit' and 'skip' parameters have been renamed. "
+                        "Please update your API client to the newest version.",
+                    )
+                else:
+                    raise HTTPException(
+                        status_code=400,
+                        detail="Request contains query parameters which cannot be used: {}".format(
+                            ", ".join(overlap)
+                        ),
+                    )
+            query: dict[Any, Any] = merge_atlas_queries(list(queries.values()))  # type: ignore
+
+            try:
+                pipeline = generate_atlas_search_pipeline(query)
+                cursor = await self.collection.aggregate(pipeline)
+                data = await cursor.to_list()
+            except (NetworkTimeout, PyMongoError) as e:
+                raise HTTPException(
+                    status_code=504 if e.timeout else 500,
+                    detail=f"Server error: {e}",
+                )
+
+            # results are returned reversed when paginating backwards so we need to fix that
+            reverse = any(
+                "$search" in p and "searchBefore" in p["$search"] for p in pipeline
+            )
+            if reverse:
+                data = list(reversed(data))
+
+            operator_meta = {}
+
+            for operator in self.query_operators:
+                data = operator.post_process(data, query)
+                operator_meta.update(operator.meta())
+
+            if data and "meta" in data[0] and data[0]["meta"]:
+                meta = Meta(
+                    total_doc=data[0]["meta"].get("count", {}).get("total", 1),
+                    facet=data[0]["meta"].get("facet", {}),
+                    pagination_token=data[-1].get("meta_pagination_token", ""),
+                )
+            else:
+                meta = Meta(total_doc=0)
+
+            response = {"data": data if data else [], "meta": {**meta.dict(), **operator_meta}}  # type: ignore
+
+            if self.disable_validation:
+                response = Response(orjson.dumps(response, default=serialization_helper))  # type: ignore
+
+            if self.header_processor is not None:
+                if self.disable_validation:
+                    self.header_processor.process_header(response, request)
+                else:
+                    self.header_processor.process_header(temp_response, request)
+
+            return response
+
+        self.router.get(
+            self.sub_path,
+            tags=self.tags,
+            summary=f"Get {model_name} documents",
+            response_model=self.response_model,
+            response_description=f"Search for a {model_name}",
+            response_model_exclude_unset=True,
+        )(attach_query_ops(search, self.query_operators))

{emmet_api-0.87.0.dev10 → emmet_api-0.87.0.dev11}/emmet/api/resource/utils.py

@@ -6,8 +6,6 @@ from emmet.api.query_operator import QueryOperator
 from emmet.api.utils import STORE_PARAMS, attach_signature
 from pymongo.asynchronous.collection import AsyncCollection
 
-NON_STORED_SOURCES = ["calcs_reversed", "orig_inputs"]
-
 
 class CollectionWithKey:
 
@@ -74,62 +72,62 @@ def generate_query_pipeline(query: dict):
 def generate_atlas_search_pipeline(query: dict):
     """
     Generate the aggregation pipeline for Atlas Search queries.
-
-    Args:
-        query: Query parameters
-        store: Store containing endpoint data
     """
     pipeline = []
+    operator: dict = {}
 
-    # generate the operator, if more than one
-    operator = {
-        "compound": {
-            "must": [q for q in query["criteria"] if not q.get("mustNot", False)]
+    if not query.get("criteria"):
+        operator = {"exists": {"path": "_id"}}
+    else:
+        operator = {"compound": {"must": [], "mustNot": []}}
+
+        # Build the must clauses
+        for q in query["criteria"]:
+            if not q.get("mustNot", False):
+                if "must" in q:
+                    # If q has a "must" key, expand its contents instead of using q directly
+                    operator["compound"]["must"].extend(q["must"])
+                else:
+                    # Use the query as-is
+                    operator["compound"]["must"].append(q)
+            else:
+                operator["compound"]["mustNot"].extend(q["mustNot"])
+
+    search_base = {
+        "$search": {
+            "index": "default",
+            "returnStoredSource": True,
+            "count": {"type": "total"},
         }
     }
-    # append the mustNot criteria to the compound operator
-    operator["compound"]["mustNot"] = [
-        q["mustNot"] for q in query["criteria"] if q.get("mustNot", False)
-    ]
-
-    if query.get("facets", False):
-        pipeline.append(
-            {
-                "$search": {
-                    "index": "default",
-                    "facet": {"operator": operator, "facets": query["facets"]},
-                }
-            }
-        )
-    else:
-        pipeline.append({"$search": {"index": "default", **operator}})
-    # add returnedStoredSource: True if non-stored source are not present in "properties"
-    # for quicker document retrieval, otherwise, do a full lookup
-    return_stored_source = not any(
-        prop in NON_STORED_SOURCES for prop in query.get("properties", [])
-    )
-    if return_stored_source:
-        pipeline[0]["$search"]["returnStoredSource"] = True  # type: ignore
 
-    sorting = query.get("sort", False)
-    if sorting:
-        # no $ sign for atlas search
-        sort_dict = {"sort": {}}  # type: ignore
-        sort_dict["sort"].update(query["sort"])
-        # add sort to $search stage
-        pipeline[0]["$search"].update(sort_dict)
+    if p_token := query.get("pagination_token", None):
+        search_base["$search"][
+            f"search{'After' if query.get('forward', True) else 'Before'}"
+        ] = p_token
+
+    search_base["$search"].update(operator)
+
+    sort_dict = {"sort": {"_id": 1}}
+    if query.get("sort", False):
+        sort_dict["sort"] = {**sort_dict["sort"], **query["sort"]}
+    # add sort to $search stage
+    search_base["$search"].update(sort_dict)
 
-    projection_dict = {"_id": 0}
+    pipeline.append(search_base)
+
+    projection_dict = {
+        "_id": 0,
+        "meta": "$$SEARCH_META",
+        "meta_pagination_token": {"$meta": "searchSequenceToken"},
+    }
     if query.get("properties", False):
         projection_dict.update({p: 1 for p in query["properties"]})
-    pipeline.insert(1, {"$project": projection_dict})  # type: ignore
+    pipeline.append({"$project": projection_dict})  # type: ignore
 
     pipeline.append({"$skip": query.get("skip", 0)})
 
     if query.get("limit", False):
         pipeline.append({"$limit": query["limit"]})
 
-    if query.get("facets", False):
-        pipeline.append({"$facet": {"docs": [], "meta": [{"$replaceWith": "$$SEARCH_META"}, {"$limit": 1}]}})  # type: ignore
-
     return pipeline

{emmet_api-0.87.0.dev10 → emmet_api-0.87.0.dev11}/emmet/api/routes/materials/electronic_structure/query_operators.py

@@ -10,7 +10,8 @@ from pymatgen.analysis.magnetism.analyzer import Ordering
 from pymatgen.core.periodic_table import Element
 from pymatgen.electronic_structure.core import OrbitalType, Spin
 
-from emmet.core.electronic_structure import BSPathType, DOSProjectionType
+from emmet.core.band_theory import BSPathType
+from emmet.core.electronic_structure import DOSProjectionType
 from emmet.core.mpid import MPID, AlphaID
 
 

emmet_api-0.87.0.dev11/emmet/api/routes/materials/materials/utils.py

@@ -0,0 +1,271 @@
+from typing import Any
+from fastapi import HTTPException
+from pymatgen.core import Composition
+from pymatgen.core.periodic_table import DummySpecies
+
+
+def formula_to_criteria(formulas: str) -> dict:
+    """
+    Santizes formula into a dictionary to search with wild cards
+
+    Arguments:
+        formula: formula with wildcards in it for unknown elements
+
+    Returns:
+        Mongo style search criteria for this formula
+    """
+    dummies = "AEGJLMQRXZ"
+
+    formula_list = [formula.strip() for formula in formulas.split(",")]
+
+    if "*" in formulas:
+        if len(formula_list) > 1:
+            raise HTTPException(
+                status_code=400,
+                detail="Wild cards only supported for single formula queries.",
+            )
+        else:
+            # Wild card in formula
+            nstars = formulas.count("*")
+
+            formula_dummies = formulas.replace("*", "{}").format(*dummies[:nstars])
+
+            try:
+                integer_formula = Composition(
+                    formula_dummies
+                ).get_integer_formula_and_factor()[0]
+            except ValueError:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Problem processing formula with wild cards.",
+                )
+
+            comp = Composition(integer_formula).reduced_composition
+            crit: dict[str, Any] = {"formula_anonymous": comp.anonymized_formula}
+            real_elts = [
+                e.name for e in comp.elements if not isinstance(e, DummySpecies)
+            ]
+            crit.update(
+                {
+                    f"composition_reduced.{el}": n
+                    for el, n in comp.to_reduced_dict.items()
+                    if el in real_elts
+                }
+            )
+            return crit
+
+    else:
+        try:
+            composition_list = [Composition(formula) for formula in formula_list]
+        except ValueError:
+            raise HTTPException(
+                status_code=400,
+                detail="Problem processing one or more provided formulas.",
+            )
+
+        if any(
+            isinstance(el, DummySpecies) for comp in composition_list for el in comp
+        ):
+            # Assume fully anonymized formula
+            if len(formula_list) == 1:
+                return {"formula_anonymous": composition_list[0].anonymized_formula}
+            else:
+                return {
+                    "formula_anonymous": {
+                        "$in": [comp.anonymized_formula for comp in composition_list]
+                    }
+                }
+
+        else:
+            if len(formula_list) == 1:
+                comp = composition_list[0]
+                # Paranoia below about floating-point "equality"
+                crit = {"nelements": len(comp)}
+
+                try:
+                    crit.update(
+                        {
+                            f"composition_reduced.{el}": n
+                            for el, n in comp.to_reduced_dict.items()
+                        }
+                    )
+                except IndexError:
+                    raise HTTPException(
+                        status_code=400,
+                        detail="Problem processing one or more provided formulas.",
+                    )
+
+                return crit
+            else:
+                return {
+                    "formula_pretty": {
+                        "$in": [comp.reduced_formula for comp in composition_list]
+                    }
+                }
+
+
+def formula_to_atlas_criteria(formulas: str) -> dict:
+    """
+    Converts formula into Atlas Search query criteria for composition-based searches
+
+    Arguments:
+        formula: formula with wildcards in it for unknown elements
+
+    Returns:
+        Atlas Search style query criteria for this formula
+    """
+    dummies = "AEGJLMQRXZ"
+    formula_list = [formula.strip() for formula in formulas.split(",")]
+    must_clauses: list[dict[str, Any]] = []
+
+    if "*" in formulas:
+        if len(formula_list) > 1:
+            raise HTTPException(
+                status_code=400,
+                detail="Wild cards only supported for single formula queries.",
+            )
+        else:
+            # Wild card in formula
+            nstars = formulas.count("*")
+            formula_dummies = formulas.replace("*", "{}").format(*dummies[:nstars])
+
+            try:
+                integer_formula = Composition(
+                    formula_dummies
+                ).get_integer_formula_and_factor()[0]
+            except ValueError:
+                raise HTTPException(
+                    status_code=400,
+                    detail="Problem processing formula with wild cards.",
+                )
+
+            comp = Composition(integer_formula).reduced_composition
+            real_elts = [
+                e.name for e in comp.elements if not isinstance(e, DummySpecies)
+            ]
+
+            # Build Atlas Search compound query
+            must_clauses = [
+                {
+                    "text": {
+                        "query": comp.anonymized_formula,
+                        "path": "formula_anonymous",
+                    }
+                }
+            ]
+
+            # Add element-specific clauses
+            must_clauses += [
+                {"equals": {"path": f"composition_reduced.{el}", "value": n}}
+                for el, n in comp.to_reduced_dict.items()
+                if el in real_elts
+            ]
+
+            return {"must": must_clauses}
+
+    else:
+        try:
+            composition_list = [Composition(formula) for formula in formula_list]
+        except ValueError:
+            raise HTTPException(
+                status_code=400,
+                detail="Problem processing one or more provided formulas.",
+            )
+
+        if any(
+            isinstance(el, DummySpecies) for comp in composition_list for el in comp
+        ):
+            # Assume fully anonymized formula
+            anonymized_formulas = [comp.anonymized_formula for comp in composition_list]
+
+            if len(formula_list) == 1:
+                return {
+                    "text": {
+                        "query": anonymized_formulas[0],
+                        "path": "formula_anonymous",
+                    }
+                }
+            else:
+                return {
+                    "in": {
+                        "path": "formula_anonymous",
+                        "value": anonymized_formulas,
+                    }
+                }
+
+        else:
+            if len(formula_list) == 1:
+                comp = composition_list[0]
+
+                # Build compound query for exact composition match
+                must_clauses = [{"equals": {"path": "nelements", "value": len(comp)}}]
+
+                try:
+                    must_clauses += [
+                        {
+                            "equals": {
+                                "path": f"composition_reduced.{el}",
+                                "value": n,
+                            }
+                        }
+                        for el, n in comp.to_reduced_dict.items()
+                    ]
+                except IndexError:
+                    raise HTTPException(
+                        status_code=400,
+                        detail="Problem processing one or more provided formulas.",
+                    )
+
+                return {"must": must_clauses}
+            else:
+                # Multiple formulas - use pretty formula matching
+                pretty_formulas = [comp.reduced_formula for comp in composition_list]
+                return {"in": {"path": "formula_pretty", "value": pretty_formulas}}
+
+
+def chemsys_to_criteria(chemsys: str) -> dict:
+    """
+    Santizes chemsys into a dictionary to search with wild cards
+
+    Arguments:
+        chemsys: A comma delimited string list of chemical systems
+            with wildcards in it for unknown elements
+
+    Returns:
+        Mongo style search criteria for this formula
+    """
+
+    crit = {}  # type: dict
+
+    chemsys_list = [chemsys_val.strip() for chemsys_val in chemsys.split(",")]
+
+    if "*" in chemsys:
+        if len(chemsys_list) > 1:
+            raise HTTPException(
+                status_code=400,
+                detail="Wild cards only supported for single chemsys queries.",
+            )
+        else:
+            eles = chemsys_list[0].split("-")
+
+            crit["nelements"] = len(eles)
+            crit.update(
+                {
+                    f"composition_reduced.{el}": {"$exists": True}
+                    for el in eles
+                    if el != "*"
+                }
+            )
+
+            return crit
+    else:
+        query_vals = [
+            "-".join(sorted(chemsys_val.split("-"))) for chemsys_val in chemsys_list
+        ]
+
+        if len(query_vals) == 1:
+            crit["chemsys"] = query_vals[0]
+        else:
+            crit["chemsys"] = {"$in": query_vals}
+
+        return crit