From 9a7278bf1044691be69190074ba6b5e0b84bd532 Mon Sep 17 00:00:00 2001 From: Tyler Mathis <35553152+tsmathis@users.noreply.github.com> Date: Mon, 10 Aug 2026 18:35:56 -0700 Subject: [PATCH 1/3] update nomad integration --- mp_api/client/mprester.py | 56 ++++++++++++++++++++++++----------- tests/client/test_mprester.py | 34 ++++++++++++++++----- 2 files changed, 65 insertions(+), 25 deletions(-) diff --git a/mp_api/client/mprester.py b/mp_api/client/mprester.py index 551f6852..e9dc8ad1 100644 --- a/mp_api/client/mprester.py +++ b/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,45 +1468,64 @@ 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 = [ + (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/") 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 @@ -1511,7 +1533,7 @@ 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): diff --git a/tests/client/test_mprester.py b/tests/client/test_mprester.py index f137b4c5..95d03f9f 100644 --- a/tests/client/test_mprester.py +++ b/tests/client/test_mprester.py @@ -727,6 +727,10 @@ 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", + ) def test_nomad_integration(self, mpr): # No particular reason for this MPID other than that it exists in NOMAD. target_mpid = "mp-10018" @@ -739,30 +743,44 @@ 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"] + ) + assert all( + 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/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( + target_mpid, file_patterns=["POSCAR", "OUTCAR"] ) 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 "re_pattern=POSCAR%7COUTCAR" in url for url in nomad_urls ) calc_type_map, nomad_urls = mpr.get_download_info( - [MPID(target_mpid)], calc_types=["GGA Deformation"] + [MPID(target_mpid)], 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 entry["calc_type"].value == "GGA Static" 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=" ) for url in nomad_urls ) From 5edb635b0801e59bd1348ebd946eb1ae8a16f3a7 Mon Sep 17 00:00:00 2001 From: Tyler Mathis <35553152+tsmathis@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:38:47 -0700 Subject: [PATCH 2/3] mypy --- mp_api/client/mprester.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mp_api/client/mprester.py b/mp_api/client/mprester.py index e9dc8ad1..3a9db361 100644 --- a/mp_api/client/mprester.py +++ b/mp_api/client/mprester.py @@ -1488,7 +1488,7 @@ def get_download_info( 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 = [ + task_id_query_params: list[tuple[str, dict]] = [ (tid, {"json_query": json.dumps({"external_id": tid})}) for tid in task_ids ] @@ -1521,7 +1521,9 @@ def get_download_info( return meta, urls - def _check_get_download_info_url_by_task_id(self, prefix, task_ids) -> list[str]: + def _check_get_download_info_url_by_task_id( + self, prefix, task_ids + ) -> list[tuple[str, dict]]: return [ pair for pair in task_ids From 2143fd4998bfba167b96fec316ab3603d4b43f3b Mon Sep 17 00:00:00 2001 From: Tyler Mathis <35553152+tsmathis@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:57:26 -0700 Subject: [PATCH 3/3] handle flaky upstream better --- tests/client/test_mprester.py | 28 ++++++++++++---------------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/tests/client/test_mprester.py b/tests/client/test_mprester.py index 95d03f9f..e1abd9b8 100644 --- a/tests/client/test_mprester.py +++ b/tests/client/test_mprester.py @@ -731,6 +731,11 @@ def test_oxygen_evolution_bad_input(self, mpr): 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" @@ -743,7 +748,9 @@ def test_nomad_integration(self, mpr): ), ): calc_type_map, nomad_urls = mpr.get_download_info( - target_mpid, file_patterns=["POSCAR"] + target_mpid, + file_patterns=["POSCAR.gz"], + calc_types=["GGA Static"], ) assert all( isinstance(entry["task_id_as_alpha"], AlphaID) @@ -757,31 +764,20 @@ def test_nomad_integration(self, mpr): ) calc_type_map, nomad_urls = mpr.get_download_info( - target_mpid, file_patterns=["POSCAR", "OUTCAR"] + target_mpid, + file_patterns=["POSCAR", "OUTCAR"], + calc_types=["GGA Static"], ) assert all( 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/v1/api/v1/entries/raw") - and "re_pattern=POSCAR%7COUTCAR" in url - for url in nomad_urls - ) - - calc_type_map, nomad_urls = mpr.get_download_info( - [MPID(target_mpid)], calc_types=["GGA Static"] - ) - assert all( - isinstance(entry["task_id_as_alpha"], AlphaID) - and entry["calc_type"].value == "GGA Static" - for entry in calc_type_map[target_mpid] - ) assert all( url.startswith( "https://nomad-lab.eu/prod/v1/api/v1/entries/raw?json_query=" ) + and "re_pattern=POSCAR%7COUTCAR" in url for url in nomad_urls )