MAPISettings added to base resource clas

materialsproject · Apr 30, 2021 · 139e9ff · 139e9ff
1 parent 7c612f6
commit 139e9ff
Showing 1 changed file with 7 additions and 8 deletions.
diff --git a/src/mp_api/core/resource.py b/src/mp_api/core/resource.py
@@ -16,6 +16,9 @@
     attach_signature,
     dynamic_import,
 )
+
+from mp_api.core.settings import MAPISettings
+
 from mp_api.core.query_operator import (
     QueryOperator,
     PaginationQuery,
@@ -307,7 +310,7 @@ async def get_by_key_versioned(
                 if version is not None:
                     version = version.replace(".", "_")
                 else:
-                    version = os.environ.get("DB_VERSION")
+                    version = os.environ.get("DB_VERSION", MAPISettings().db_version)
 
                 prefix = self.store.collection_name.split("_")[0]
                 self.store.collection_name = f"{prefix}_{version}"
@@ -316,9 +319,7 @@ async def get_by_key_versioned(
 
                 crit = {self.store.key: key}
 
-                if model_name == "MaterialsCoreDoc":
-                    crit.update({"_sbxn": "core"})
-                elif model_name == "TaskDoc":
+                if model_name == "TaskDoc":
                     crit.update({"sbxn": "core"})
 
                 item = [
@@ -380,16 +381,14 @@ async def search(**queries: STORE_PARAMS):
                     query["criteria"].pop("version")
 
                 else:
-                    version = os.environ.get("DB_VERSION")
+                    version = os.environ.get("DB_VERSION", MAPISettings().db_version)
 
                 prefix = self.store.collection_name.split("_")[0]
                 self.store.collection_name = f"{prefix}_{version}"
 
             self.store.connect(force_reset=True)
 
-            if model_name == "MaterialsCoreDoc":
-                query["criteria"].update({"_sbxn": "core"})
-            elif model_name == "TaskDoc":
+            if model_name == "TaskDoc":
                 query["criteria"].update({"sbxn": "core"})
 
             data = list(self.store.query(**query))  # type: ignore