API to update extras field

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

Hide file tree

Showing 2 changed files with 41 additions and 27 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",)
@@ -139,37 +139,50 @@ 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_optimade.common import AiidaError
+        from aiida import orm
 
         optimade = self._get_optimade_extras()
         if optimade:
             optimade.update(self.new_attributes)
         else:
             optimade = self.new_attributes
-        extras = (
-            self._get_unique_node_property("extras")
-            if self._get_unique_node_property("extras")
-            else {}
-        )
-        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
-
-            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}"
-            )
+        # extras = (
+        #     self._get_unique_node_property("extras")
+        #     if self._get_unique_node_property("extras")
+        #     else {}
+        # )
+        # extras[self.EXTRAS_KEY] = optimade
+
+        # profile = get_manager().get_profile()
+
+        # 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
+        #     LOGGER.debug(f"{self._node.extras}")
+        #     # LOGGER.debug(f"extras: {extras}")
+
+        #     # self._node.reset_extras(extras)
+        #     node = orm.load_node(self._pk)
+        #     node.set_extra(self.EXTRAS_KEY, optimade)
+        #     LOGGER.debug(f"{self._node.extras}")
+        #     # LOGGER.debug(f"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}"
+        #     )
+
+        LOGGER.debug("Updating Node %s in AiiDA DB!", self._pk)
+        node = orm.load_node(self._pk)
+        node.set_extra(self.EXTRAS_KEY, optimade)
 
         # For posterity, this is how to do the same, going through AiiDA's API:
+        # self._node.set_extra_many(extras)
         # self._node.set_extra(self.EXTRAS_KEY, optimade)
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