API to update extras field

aiidateam · May 30, 2022 · 1883537 · 1883537
1 parent cd2cfa3
commit 1883537
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 20 deletions.
diff --git a/aiida_optimade/translators/entities.py b/aiida_optimade/translators/entities.py
@@ -2,7 +2,7 @@
 from aiida.orm.nodes import Node
 from aiida.orm.querybuilder import QueryBuilder
 
-from aiida_optimade.common import AiidaEntityNotFound, AiidaError, LOGGER
+from aiida_optimade.common import AiidaEntityNotFound, LOGGER
 
 
 __all__ = ("AiidaEntityTranslator",)
@@ -138,7 +138,7 @@ def _store_attributes_mongo(self) -> None:
 
     def _store_attributes_node_extra(self) -> None:
         """Store new attributes in Node extras"""
-        from aiida.manage.manager import get_manager
+        # from aiida.manage.manager import get_manager
 
         optimade = self._get_optimade_extras()
         if optimade:
@@ -152,24 +152,22 @@ def _store_attributes_node_extra(self) -> None:
         )
         extras[self.EXTRAS_KEY] = optimade
 
-        profile = get_manager().get_profile()
-        LOGGER.debug(
-            "(%s) Updating Node %s in AiiDA DB!", profile.storage_backend, self._pk
-        )
-        if profile.storage_backend == "psql_dos":
-            # TODO: this can be `set_extra` directly? is there performance issue?
-            from aiida.storage.psql_dos.models.node import DbNode
+        # profile = get_manager().get_profile()
+        LOGGER.debug("Updating Node %s in AiiDA DB!", self._pk)
+        # if profile.storage_backend == "psql_dos":
+        #     # TODO: this can be `set_extra` directly? is there performance issue?
+        #     from aiida.storage.psql_dos.models.node import DbNode
 
-            with get_manager().get_backend().transaction() as session:
-                session.query(DbNode).filter(DbNode.id == self._pk).update(
-                    values={"extras": extras}
-                )
+        #     with get_manager().get_backend().transaction() as session:
+        #         session.query(DbNode).filter(DbNode.id == self._pk).update(
+        #             values={"extras": extras}
+        #         )
 
-        else:
-            raise AiidaError(
-                f'Unknown AiiDA backend "{profile.database_backend}" for profile'
-                f"{profile}"
-            )
+        # else:
+        #     raise AiidaError(
+        #         f'Unknown AiiDA backend "{profile.database_backend}" for profile'
+        #         f"{profile}"
+        #     )
 
         # For posterity, this is how to do the same, going through AiiDA's API:
-        # self._node.set_extra(self.EXTRAS_KEY, optimade)
+        self._node.set_extra_many(extras)
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,5 @@
-aiida-core[atomic_tools]~=2.0.1
+aiida-core~=2.0.1
+pymatgen>=2019.7.2,<=2022.1.9,!=2019.9.7
 optimade[mongo]~=0.17.2
 uvicorn~=0.17.6