allow for version setting via client

Author: esoteric-ephemera

mp_api/client/core/client.py

@@ -134,6 +134,7 @@ def __init__(
         session: requests.Session | None = None,
         headers: dict | None = None,
         mute_progress_bars: bool = MAPI_CLIENT_SETTINGS.MUTE_PROGRESS_BARS,
+        db_version: str | None = None,
         local_dataset_cache: (
             str | os.PathLike
         ) = MAPI_CLIENT_SETTINGS.LOCAL_DATASET_CACHE,
@@ -167,6 +168,9 @@ def __init__(
                 advanced usage only.
             headers: Custom headers for localhost connections.
             mute_progress_bars: Whether to disable progress bars.
+            db_version (str) : EXPERIMENTAL, allows for accessing a different version of the database
+                than what is currently deployed. The Materials Project cannot guarantee that all
+                features will still work.
             local_dataset_cache: Target directory for downloading full datasets. Defaults
                 to 'mp_datasets' in the user's home directory
             force_renew: Option to overwrite existing local dataset
@@ -192,6 +196,7 @@ def __init__(
 
         self.use_document_model = use_document_model
         self.mute_progress_bars = mute_progress_bars
+        self.db_version: str = db_version or ""
         self.local_dataset_cache = Path(local_dataset_cache)
         self.force_renew = force_renew
         self._query_builder = (
@@ -260,6 +265,45 @@ def __exit__(self, exc_type, exc_val, exc_tb):  # pragma: no cover
             self.session.close()
         self._session = None
 
+    @staticmethod
+    @cache
+    def _get_heartbeat_info(endpoint) -> tuple[str, list[str]]:
+        """DB version:
+        The Materials Project database is periodically updated and has a
+        database version associated with it. When the database is updated,
+        consolidated data (information about "a material") may and does
+        change, while calculation data about a specific calculation task
+        remains unchanged and available for querying via its task_id.
+
+        The database version is set as a date in the format YYYY_MM_DD,
+        where "_DD" may be optional. An additional numerical or `postN` suffix
+        might be added if multiple releases happen on the same day.
+
+        Access Controlled Datasets:
+        Certain contributions to the Materials Project have access
+        control restrictions that require explicit agreement to the
+        Terms of Use for the respective datasets prior to access being
+        granted.
+
+        A full list of the Terms of Use for all contributions in the
+        Materials Project are available at:
+
+        https://next-gen.materialsproject.org/about/terms
+
+        Returns:
+            tuple with database version as a string and a comma separated
+            string with all calculation batch identifiers that have access
+            restrictions
+        """
+        if (get_resp := requests.get(url=endpoint + "heartbeat")).status_code == 403:
+            _emit_status_warning()
+            return (
+                "",
+                [],
+            )  # Catiously do not allow access to any access controlled `batch_id`s
+        response = get_resp.json()
+        return response["db_version"], response["access_controlled_batch_ids"]
+
 
 class BaseRester(_Rester):
     """Base client class with core stubs."""
@@ -278,6 +322,7 @@ def __init__(
         session: requests.Session | None = None,
         headers: dict | None = None,
         mute_progress_bars: bool = MAPI_CLIENT_SETTINGS.MUTE_PROGRESS_BARS,
+        db_version: str | None = None,
         local_dataset_cache: (
             str | os.PathLike
         ) = MAPI_CLIENT_SETTINGS.LOCAL_DATASET_CACHE,
@@ -316,6 +361,9 @@ def __init__(
                 and will not give auto-complete for available fields.
             headers: Custom headers for localhost connections.
             mute_progress_bars: Whether to disable progress bars.
+            db_version (str) : EXPERIMENTAL, allows for accessing a different version of the database
+                than what is currently deployed. The Materials Project cannot guarantee that all
+                features will still work.
             local_dataset_cache: Target directory for downloading full datasets. Defaults
                 to 'mp_datasets' in the user's home directory
             force_renew: Option to overwrite existing local dataset
@@ -333,6 +381,7 @@ def __init__(
             session=session,
             headers=headers,
             mute_progress_bars=mute_progress_bars,
+            db_version=db_version,
             local_dataset_cache=local_dataset_cache,
             force_renew=force_renew,
             query_builder=query_builder,
@@ -343,9 +392,11 @@ def __init__(
         self.endpoint = validate_endpoint(endpoint, suffix=self.suffix)
 
         (
-            self.db_version,
+            hb_db_version,
             self.access_controlled_batch_ids,
-        ) = BaseRester._get_heartbeat_info(self.base_endpoint)
+        ) = self._get_heartbeat_info(self.base_endpoint)
+        if not self.db_version:
+            self.db_version = hb_db_version
 
         self.timeout = timeout
         self._s3_client = s3_client
@@ -359,45 +410,6 @@ def s3_client(self):
             )
         return self._s3_client
 
-    @staticmethod
-    @cache
-    def _get_heartbeat_info(endpoint) -> tuple[str, list[str]]:
-        """DB version:
-        The Materials Project database is periodically updated and has a
-        database version associated with it. When the database is updated,
-        consolidated data (information about "a material") may and does
-        change, while calculation data about a specific calculation task
-        remains unchanged and available for querying via its task_id.
-
-        The database version is set as a date in the format YYYY_MM_DD,
-        where "_DD" may be optional. An additional numerical or `postN` suffix
-        might be added if multiple releases happen on the same day.
-
-        Access Controlled Datasets:
-        Certain contributions to the Materials Project have access
-        control restrictions that require explicit agreement to the
-        Terms of Use for the respective datasets prior to access being
-        granted.
-
-        A full list of the Terms of Use for all contributions in the
-        Materials Project are available at:
-
-        https://next-gen.materialsproject.org/about/terms
-
-        Returns:
-            tuple with database version as a string and a comma separated
-            string with all calculation batch identifiers that have access
-            restrictions
-        """
-        if (get_resp := requests.get(url=endpoint + "heartbeat")).status_code == 403:
-            _emit_status_warning()
-            return (
-                "",
-                [],
-            )  # Catiously do not allow access to any access controlled `batch_id`s
-        response = get_resp.json()
-        return response["db_version"], response["access_controlled_batch_ids"]
-
     def _post_resource(
         self,
         body: dict | None = None,
@@ -744,17 +756,23 @@ def _query_delta_backed(
 
         predicate = (
             f"WHERE batch_id NOT IN ({controlled_batch_str})"
-            if not has_gnome_access
+            if not has_gnome_access and controlled_batch_str
             else ""
         )
+        # TODO: do we need something like this?
+        # predicate += f"{' AND ' if predicate else 'WHERE '}version='{self.db_version}'"
 
         # Setup progress bar
         num_docs_needed: int = tbl.count()
 
         if not has_gnome_access:
-            num_docs_needed = self.count(
-                {"batch_id_neq_any": self.access_controlled_batch_ids}
-            )
+            try:
+                num_docs_needed = self.count(
+                    {"batch_id_neq_any": self.access_controlled_batch_ids}
+                )
+            except MPRestError:
+                # batch_id isn't a valid field
+                num_docs_needed = self.count()
 
         pbar = (
             tqdm(
@@ -918,15 +936,19 @@ def _query_resource(
                 elif suffix in STATIC_COLLECTIONS:
                     bucket_suffix = "build"
                     prefix = f"static-collections/{suffix}"
+                elif self.delta_backed:
+                    return self._query_delta_backed(
+                        "materialsproject-build",
+                        f"collections/{suffix}/",
+                        timeout=timeout,
+                    )
                 else:
+                    # TODO: remove once all collections are migrated to delta-backed format
                     bucket_suffix = "build"
                     prefix = f"collections/{self.db_version.replace('.', '-')}/{suffix}"
 
                 bucket = f"materialsproject-{bucket_suffix}"
 
-                if self.delta_backed:
-                    return self._query_delta_backed(bucket, prefix, timeout=timeout)
-
                 # Paginate over all entries in the bucket.
                 # TODO: change when a subset of entries needed from DB
                 paginator = self.s3_client.get_paginator("list_objects_v2")
@@ -1671,6 +1693,7 @@ def __getattr__(self, v: str):
                     use_document_model=self.use_document_model,
                     headers=self.headers,
                     mute_progress_bars=self.mute_progress_bars,
+                    db_version=self.db_version,
                     local_dataset_cache=self.local_dataset_cache,
                     force_renew=self.force_renew,
                     query_builder=self._query_builder,

mp_api/client/core/settings.py

@@ -15,7 +15,7 @@
 _MAX_LIST_LENGTH = min(PMG_SETTINGS.get("MPRESTER_MAX_LIST_LENGTH", 10000), 10000)
 
 _EMMET_SETTINGS = EmmetSettings()  # type: ignore[call-arg]
-_DEFAULT_ENDPOINT = "https://api.materialsproject.org/"
+_DEFAULT_ENDPOINT = "http://0.0.0.0:8000"  # https://api.materialsproject.org/"
 
 
 class MAPIClientSettings(BaseSettings):

mp_api/client/mprester.py

@@ -97,6 +97,7 @@ def __init__(
         session: Session | None = None,
         headers: dict | None = None,
         mute_progress_bars: bool = MAPI_CLIENT_SETTINGS.MUTE_PROGRESS_BARS,
+        db_version: str | None = None,
         local_dataset_cache: (
             str | os.PathLike
         ) = MAPI_CLIENT_SETTINGS.LOCAL_DATASET_CACHE,
@@ -130,6 +131,9 @@ def __init__(
             session: Session object to use. By default (None), the client will create one.
             headers: Custom headers for localhost connections.
             mute_progress_bars:  Whether to mute progress bars.
+            db_version (str) : EXPERIMENTAL, allows for accessing a different version of the database
+                than what is currently deployed. The Materials Project cannot guarantee that all
+                features will still work.
             local_dataset_cache: Target directory for downloading full datasets. Defaults
                 to "mp_datasets" in the user's home directory
             force_renew: Option to overwrite existing local dataset
@@ -152,6 +156,7 @@ def __init__(
             session=session,
             headers=headers,
             mute_progress_bars=mute_progress_bars,
+            db_version=db_version,
             local_dataset_cache=local_dataset_cache,
             force_renew=force_renew,
             query_builder=query_builder,
@@ -210,6 +215,17 @@ def __init__(
                 stacklevel=2,
             )
 
+        if self.db_version:
+            warnings.warn(
+                "Specifying an explicit database version is an experimental "
+                "feature. The Materials Project cannot guarantee "
+                "functionality at this time, use at your own risk!",
+                stacklevel=2,
+                category=MPRestWarning,
+            )
+        else:
+            self.db_version = self._get_heartbeat_info(self.endpoint)[0]
+
         if notify_db_version:
             self._db_version_check()
 
@@ -233,6 +249,7 @@ def __init__(
                         use_document_model=self.use_document_model,
                         headers=self.headers,
                         mute_progress_bars=self.mute_progress_bars,
+                        db_version=self.db_version,
                         local_dataset_cache=self.local_dataset_cache,
                         force_renew=self.force_renew,
                         query_builder=self._query_builder,
@@ -300,8 +317,7 @@ def __dir__(self):
         )
 
     def __repr__(self) -> str:
-        db_version = self.get_database_version()
-        return f"MPRester({'v' + db_version if db_version else 'unknown version'})"
+        return f"MPRester({'v' + self.self.db_version if self.db_version else 'unknown version'})"
 
     def get_task_ids_associated_with_material_id(
         self, material_id: str, calc_types: list[CalcType] | None = None
@@ -364,7 +380,9 @@ def get_structure_by_material_id(
         return structure_data
 
     def get_database_version(self) -> str | None:
-        """The Materials Project database is periodically updated and has a
+        """DEPRECATED: use `self.db_version` instead.
+
+        The Materials Project database is periodically updated and has a
         database version associated with it. When the database is updated,
         consolidated data (information about "a material") may and does
         change, while calculation data about a specific calculation task
@@ -376,10 +394,13 @@ def get_database_version(self) -> str | None:
 
         Returns: database version as a string if accessible, None otherwise
         """
-        if (get_resp := get(url=self.endpoint + "heartbeat")).status_code == 403:
-            _emit_status_warning()
-            return None
-        return get_resp.json()["db_version"]
+        warnings.warn(
+            "`get_database_version` has been deprecated in favor of "
+            "MPRester().db_version.",
+            stacklevel=2,
+            category=MPRestWarning,
+        )
+        return self.db_version
 
     @staticmethod
     @cache
@@ -1676,7 +1697,6 @@ def _db_version_check(self) -> None:
         """Check if the database version has drifted."""
         import yaml  # type: ignore[import-untyped]
 
-        db_version = self.get_database_version()
         old_db_version = None
         if MAPI_CLIENT_SETTINGS.LOG_FILE.exists():
             old_db_version = (
@@ -1687,15 +1707,15 @@ def _db_version_check(self) -> None:
             if not isinstance(old_db_version, str):
                 old_db_version = None
 
-        if old_db_version != db_version:
+        if old_db_version != self.db_version:
             MAPI_CLIENT_SETTINGS.LOG_FILE.write_text(
-                yaml.safe_dump({"MAPI_DB_VERSION": db_version})
+                yaml.safe_dump({"MAPI_DB_VERSION": self.db_version})
             )
 
             if old_db_version:
                 warnings.warn(
                     "Materials Project database version has changed "
-                    f"from v{old_db_version} to v{db_version}.",
+                    f"from v{old_db_version} to v{self.db_version}.",
                     category=MPRestWarning,
                     stacklevel=2,
                 )

tests/client/materials/test_thermo.py

@@ -64,7 +64,8 @@ def test_client(rester):
 def test_get_phase_diagram_from_chemsys():
     # Test that a phase diagram is returned
 
+    pd = ThermoRester().get_phase_diagram_from_chemsys("Hf-Pm", thermo_type="GGA_GGA+U")
     assert isinstance(
-        ThermoRester().get_phase_diagram_from_chemsys("Hf-Pm", thermo_type="GGA_GGA+U"),
+        pd,
         PhaseDiagram,
     )

tests/client/test_mprester.py

@@ -54,7 +54,9 @@ def mpr():
 @requires_api_key
 class TestMPRester:
     fake_mp_api_key = "12345678901234567890123456789012"  # 32 chars
-    default_endpoint = _DEFAULT_ENDPOINT
+    default_endpoint = _DEFAULT_ENDPOINT + (
+        "/" if not _DEFAULT_ENDPOINT.endswith("/") else ""
+    )
 
     def test_get_structure_by_material_id(self, mpr):
         s0 = mpr.get_structure_by_material_id("mp-149")
@@ -67,9 +69,14 @@ def test_get_structure_by_material_id(self, mpr):
         assert {s.formula for s in s2} == {"Si2"}
 
     def test_get_database_version(self, mpr):
-        db_version = mpr.get_database_version()
+        db_version = mpr.db_version
         assert db_version is not None
 
+        with pytest.warns(
+            MPRestWarning, match="`get_database_version` has been deprecated"
+        ):
+            assert db_version == mpr.get_database_version()
+
     def test_get_material_id_from_task_id(self, mpr):
         assert mpr.get_material_id_from_task_id("mp-540081") == "mp-19017"
 
@@ -710,7 +717,7 @@ def test_db_warning(self, monkeypatch: pytest.MonkeyPatch):
             monkeypatch.setattr(MAPI_CLIENT_SETTINGS, "LOG_FILE", Path(tmp_log.name))
 
             with MPRester(notify_db_version=True) as mpr:
-                db_version = mpr.get_database_version()
+                db_version = mpr.db_version
 
             parsed_db_ver = yaml.safe_load(Path(tmp_log.name).read_text()).get(
                 "MAPI_DB_VERSION"