migrate web utils

Author: esoteric-ephemera

mp_api/client/_server_utils.py

@@ -0,0 +1,103 @@
+"""Define utilities needed by the MP web server."""
+from __future__ import annotations
+
+try:
+    import flask
+except ImportError:
+    from mp_api.client.core.exceptions import MPRestError
+
+    raise MPRestError("`flask` must be installed to use server utilities.")
+
+import requests
+
+from mp_api.client import MPRester
+from mp_api.client.core.utils import validate_api_key
+
+SESSION = requests.Session()
+
+
+def is_localhost() -> bool:
+    """Determine if current env is local or production.
+
+    Returns:
+        bool: True if the environment is locally hosted.
+    """
+    return (
+        True
+        if not flask.has_request_context()
+        else flask.request.headers.get("Host", "").startswith(
+            ("localhost:", "127.0.0.1:", "0.0.0.0:")
+        )
+    )
+
+
+def get_consumer() -> dict[str, str]:
+    """Identify the consumer associated with the current request.
+
+    Returns:
+        dict of str to str, the headers associated with the consumer
+    """
+    if not flask.has_request_context():
+        return {}
+
+    names = [
+        "X-Consumer-Id",  # kong internal uuid
+        "X-Consumer-Custom-Id",  # api key
+        "X-Consumer-Username",  # <provider>:<email>
+        "X-Anonymous-Consumer",  # is anonymous user?
+        "X-Authenticated-Groups",  # groups this user belongs to
+        "X-Consumer-Groups",  # same as X-Authenticated-Groups
+    ]
+    headers = flask.request.headers
+    return {name: headers[name] for name in names if headers.get(name) is not None}
+
+
+def is_logged_in_user(consumer: dict[str, str] | None = None) -> bool:
+    """Check if the client has the necessary headers for an authenticated user.
+
+    Args:
+        consumer (dict of str to str, or None): Headers associated with the consumer
+
+    Returns:
+        bool : True if the consumer is a logged-in user, False otherwise.
+    """
+    c = consumer or get_consumer()
+    return bool(not c.get("X-Anonymous-Consumer") and c.get("X-Consumer-Id"))
+
+
+def get_user_api_key(consumer: dict[str, str] | None = None) -> str | None:
+    """Get the api key that belongs to the current user.
+
+    If running on localhost, api key is obtained from
+    the environment variable MP_API_KEY.
+
+    Args:
+        consumer (dict of str to str, or None): Headers associated with the consumer
+
+    Returns:
+        str, the API key, or None if no API key could be identified.
+    """
+    c = consumer or get_consumer()
+
+    if is_localhost():
+        return validate_api_key()
+    elif is_logged_in_user(c):
+        return c.get("X-Consumer-Custom-Id")
+    return None
+
+
+def get_rester(**kwargs) -> MPRester:
+    """Create MPRester with headers set for localhost and production compatibility.
+
+    Args:
+        **kwargs : kwargs to pass to MPRester
+
+    Returns:
+        MPRester
+    """
+    if is_localhost():
+        dev_api_key = get_user_api_key()
+        SESSION.headers["x-api-key"] = dev_api_key
+        return MPRester(api_key=dev_api_key, session=SESSION, **kwargs)
+
+    return MPRester(headers=get_consumer(), session=SESSION, **kwargs)

mp_api/client/routes/_server.py

@@ -1,3 +1,5 @@
+"""Define RESTers required by the MP web server."""
+
 from __future__ import annotations
 
 from typing import TYPE_CHECKING

pyproject.toml

@@ -34,12 +34,14 @@ dynamic = ["version"]
 
 [project.optional-dependencies]
 mcp = ["fastmcp"]
+server = ["flask"]
 all = [
-  "emmet-core[all]>=0.86.3",
   "custodian",
-  "mpcontribs-client>=5.10",
+  "emmet-core[all]>=0.86.3",
   "fastmcp",
+  "flask",
   "matminer>=0.10.0",
+  "mpcontribs-client>=5.10",
 ]
 test = [
   "pre-commit",