From c536639b9ec925dd2d0cb80fa2165b24c67eed15 Mon Sep 17 00:00:00 2001
From: Jusong Yu <jusong.yeu@gmail.com>
Date: Mon, 30 May 2022 10:38:27 +0200
Subject: [PATCH] API to update extras field

---
 aiida_optimade/translators/entities.py | 65 +++++++++++++++-----------
 requirements.txt                       |  3 +-
 2 files changed, 41 insertions(+), 27 deletions(-)

diff --git a/aiida_optimade/translators/entities.py b/aiida_optimade/translators/entities.py
index 3ff677c4..3dfc31a4 100644
--- 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
index 8dd08a1a..d5901f91 100644
--- 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