update nomad integration

mp_api/client/mprester.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
 import itertools
+import json
 import os
+import re
 import warnings
 from collections import defaultdict
 from functools import cache, lru_cache
 from typing import TYPE_CHECKING
+from urllib.parse import urlencode
 
 from emmet.core.band_theory import BSPathType
 from emmet.core.mpid import MPID, AlphaID
@@ -1465,53 +1468,74 @@ def get_download_info(
                 if calc_types and calc_type not in calc_types:
                     continue
                 mp_id = doc["material_id"]
-                meta[mp_id].append({"task_id": task_id, "calc_type": calc_type})
+                meta[mp_id].append(
+                    {
+                        "task_id": str(AlphaID(task_id, prefix="mp").formatted),
+                        "task_id_as_alpha": task_id,
+                        "calc_type": calc_type,
+                    }
+                )
 
         if not meta:
             raise ValueError(f"No tasks found for material id {material_ids}.")
 
         # return a list of URLs for NoMaD Downloads containing the list of files
         # for every external_id in `task_ids`
-        # For reference, please visit https://nomad-lab.eu/prod/rae/api/
+        # For reference, please visit https://nomad-lab.eu/prod/v1/api/v1/extensions/docs#/entries/raw
 
         # check if these task ids exist on NOMAD
-        prefix = "https://nomad-lab.eu/prod/rae/api/repo/?"
-        if file_patterns is not None:
-            for file_pattern in file_patterns:
-                prefix += f"file_pattern={file_pattern}&"
-        prefix += "external_id="
+        nomad_check_endpoint = "https://nomad-lab.eu/prod/v1/api/v1/entries/rawdir"
+        nomad_download_endpoint = "https://nomad-lab.eu/prod/v1/api/v1/entries/raw"
 
         task_ids = [t["task_id"] for tl in meta.values() for t in tl]
+        task_id_query_params: list[tuple[str, dict]] = [
+            (tid, {"json_query": json.dumps({"external_id": tid})}) for tid in task_ids
+        ]
+
         nomad_exist_task_ids = self._check_get_download_info_url_by_task_id(
-            prefix=prefix, task_ids=task_ids
+            prefix=nomad_check_endpoint, task_ids=task_id_query_params
         )
+
         if len(nomad_exist_task_ids) != len(task_ids):
             self._print_help_message(
-                nomad_exist_task_ids, task_ids, file_patterns, calc_types
+                [pair[0] for pair in nomad_exist_task_ids],
+                task_ids,
+                file_patterns,
+                calc_types,
             )
 
         # generate download links for those that exist
-        prefix = "https://nomad-lab.eu/prod/rae/api/raw/query?"
         if file_patterns is not None:
-            for file_pattern in file_patterns:
-                prefix += f"file_pattern={file_pattern}&"
-        prefix += "external_id="
+            if len(file_patterns) == 1:
+                for _, json_query in nomad_exist_task_ids:
+                    json_query["glob_pattern"] = file_patterns[0]
+            else:
+                re_pattern = "|".join(re.escape(pattern) for pattern in file_patterns)
+                for _, json_query in nomad_exist_task_ids:
+                    json_query["re_pattern"] = re_pattern
+
+        urls = [
+            f"{nomad_download_endpoint}?{urlencode(json_query)}"
+            for _, json_query in nomad_exist_task_ids
+        ]
 
-        urls = [prefix + tids for tids in nomad_exist_task_ids]
         return meta, urls
 
-    def _check_get_download_info_url_by_task_id(self, prefix, task_ids) -> list[str]:
-        prefix = prefix.replace("/raw/query", "/repo/")
+    def _check_get_download_info_url_by_task_id(
+        self, prefix, task_ids
+    ) -> list[tuple[str, dict]]:
         return [
-            task_id for task_id in task_ids if self._check_nomad_exist(prefix + task_id)
+            pair
+            for pair in task_ids
+            if self._check_nomad_exist(url=f"{prefix}?{urlencode(pair[1])}")
         ]
 
     @staticmethod
     def _check_nomad_exist(url) -> bool:
         response = get(url=url)
         if response.status_code != 200:
             return False
-        return load_json(response.text)["pagination"]["total"] != 0
+        return response.json()["pagination"]["total"] != 0
 
     @staticmethod
     def _print_help_message(nomad_exist_task_ids, task_ids, file_patterns, calc_types):

tests/client/test_mprester.py

@@ -727,6 +727,15 @@ def test_oxygen_evolution_bad_input(self, mpr):
         with pytest.raises(ValueError, match="No available insertion electrode data"):
             _ = mpr.get_oxygen_evolution("mp-2207", "Al")
 
+    @pytest.mark.skipif(
+        os.environ.get("GITHUB_ACTIONS") != "true",
+        reason="Slow - don't want to impede local dev",
+    )
+    @pytest.mark.xfail(
+        raises=requests.exceptions.ConnectionError,
+        reason="upstream known to timeout",
+        strict=False,
+    )
     def test_nomad_integration(self, mpr):
         # No particular reason for this MPID other than that it exists in NOMAD.
         target_mpid = "mp-10018"
@@ -739,31 +748,36 @@ def test_nomad_integration(self, mpr):
             ),
         ):
             calc_type_map, nomad_urls = mpr.get_download_info(
-                target_mpid, file_patterns=["some_pattern"]
+                target_mpid,
+                file_patterns=["POSCAR.gz"],
+                calc_types=["GGA Static"],
             )
             assert all(
-                isinstance(entry["task_id"], AlphaID)
+                isinstance(entry["task_id_as_alpha"], AlphaID)
                 and isinstance(entry["calc_type"], CalcType)
                 for entry in calc_type_map[target_mpid]
             )
             assert all(
-                url.startswith("https://nomad-lab.eu/prod/rae/api/raw/query")
-                and "file_pattern=some_pattern" in url
+                url.startswith("https://nomad-lab.eu/prod/v1/api/v1/entries/raw")
+                and "glob_pattern=POSCAR" in url
                 for url in nomad_urls
             )
 
             calc_type_map, nomad_urls = mpr.get_download_info(
-                [MPID(target_mpid)], calc_types=["GGA Deformation"]
+                target_mpid,
+                file_patterns=["POSCAR", "OUTCAR"],
+                calc_types=["GGA Static"],
             )
             assert all(
-                isinstance(entry["task_id"], AlphaID)
-                and entry["calc_type"].value == "GGA Deformation"
+                isinstance(entry["task_id_as_alpha"], AlphaID)
+                and isinstance(entry["calc_type"], CalcType)
                 for entry in calc_type_map[target_mpid]
             )
             assert all(
                 url.startswith(
-                    "https://nomad-lab.eu/prod/rae/api/raw/query?external_id="
+                    "https://nomad-lab.eu/prod/v1/api/v1/entries/raw?json_query="
                 )
+                and "re_pattern=POSCAR%7COUTCAR" in url
                 for url in nomad_urls
             )