diff --git a/mp_api/client/mprester.py b/mp_api/client/mprester.py index 551f6852..3a9db361 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,66 @@ 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 @@ -1511,7 +1535,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..e1abd9b8 100644 --- a/tests/client/test_mprester.py +++ b/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 )