def test_replacement_1(self):
from aiida.backends.djsite.db.models import DbExtra
n1 = Node().store()
n2 = Node().store()
DbExtra.set_value_for_node(n1._dbnode, "pippo", [1, 2, 'a'])
DbExtra.set_value_for_node(n1._dbnode, "pippobis", [5, 6, 'c'])
DbExtra.set_value_for_node(n2._dbnode, "pippo2", [3, 4, 'b'])
self.assertEquals(n1.get_extras(), {'pippo': [1, 2, 'a'],
'pippobis': [5, 6, 'c'],
'_aiida_hash': n1.get_hash()
})
self.assertEquals(n2.get_extras(), {'pippo2': [3, 4, 'b'],
'_aiida_hash': n2.get_hash()
})
new_attrs = {"newval1": "v", "newval2": [1, {"c": "d", "e": 2}]}
DbExtra.reset_values_for_node(n1._dbnode, attributes=new_attrs)
self.assertEquals(n1.get_extras(), new_attrs)
self.assertEquals(n2.get_extras(), {'pippo2': [3, 4, 'b'], '_aiida_hash': n2.get_hash()})
DbExtra.del_value_for_node(n1._dbnode, key='newval2')
del new_attrs['newval2']
self.assertEquals(n1.get_extras(), new_attrs)
# Also check that other nodes were not damaged
self.assertEquals(n2.get_extras(), {'pippo2': [3, 4, 'b'], '_aiida_hash': n2.get_hash()})
def _del_db_extra(self, key):
from aiida.backends.djsite.db.models import DbExtra
if not DbExtra.has_key(self.dbnode, key):
raise AttributeError("DbExtra {} does not exist".format(
key))
return DbExtra.del_value_for_node(self.dbnode, key)
self._increment_version_number_db()
def _set_db_extra(self, key, value, exclusive=False):
from aiida.backends.djsite.db.models import DbExtra
DbExtra.set_value_for_node(self.dbnode,
key,
value,
stop_if_existing=exclusive)
self._increment_version_number_db()
def del_extra(self, key):
from aiida.backends.djsite.db.models import DbExtra
if self._to_be_stored:
raise ModificationNotAllowed(
"The extras of a node can be set and deleted "
"only after storing the node")
if not DbExtra.has_key(self.dbnode, key):
raise AttributeError("DbExtra {} does not exist".format(key))
return DbExtra.del_value_for_node(self.dbnode, key)
self._increment_version_number_db()
def set_extra_exclusive(self, key, value):
from aiida.backends.djsite.db.models import DbExtra
DbExtra.validate_key(key)
if self._to_be_stored:
raise ModificationNotAllowed(
"The extras of a node can be set only after "
"storing the node")
DbExtra.set_value_for_node(self.dbnode,
key,
value,
stop_if_existing=True)
self._increment_version_number_db()
def extras(self):
from aiida.backends.djsite.db.models import DbExtra
if self._to_be_stored:
return
else:
extraslist = DbExtra.list_all_node_elements(self.dbnode)
for e in extraslist:
yield e.key
def iterextras(self):
from aiida.backends.djsite.db.models import DbExtra
if self._to_be_stored:
# If it is not stored yet, there are no extras that can be
# added (in particular, we do not even have an ID to use!)
# Return without value, meaning that this is an empty generator
return
else:
extraslist = DbExtra.list_all_node_elements(self.dbnode)
for e in extraslist:
yield (e.key, e.getvalue())
def get_extra(self, key, *args):
from aiida.backends.djsite.db.models import DbExtra
if len(args) > 1:
raise ValueError("After the key name you can pass at most one"
"value, that is the default value to be used "
"if no extra is found.")
try:
if self._to_be_stored:
raise AttributeError("DbExtra '{}' does not exist yet, the "
"node is not stored".format(key))
else:
return DbExtra.get_value_for_node(dbnode=self.dbnode, key=key)
except AttributeError as e:
try:
return args[0]
except IndexError:
raise e
def _get_aiida_res(self, key, res):
"""
Some instance returned by ORM (django or SA) need to be converted
to Aiida instances (eg nodes)
:param res: the result returned by the query
:param key: the key that this entry would be return with
:returns: an aiida-compatible instance
"""
if key.startswith('attributes.'):
# If you want a specific attributes, that key was stored in res.
# So I call the getvalue method to expand into a dictionary
try:
returnval = DbAttribute.objects.get(id=res).getvalue()
except ObjectDoesNotExist:
# If the object does not exist, return None. This is consistent
# with SQLAlchemy inside the JSON
returnval = None
elif key.startswith('extras.'):
# Same as attributes
try:
returnval = DbExtra.objects.get(id=res).getvalue()
except ObjectDoesNotExist:
returnval = None
elif key == 'attributes':
# If you asked for all attributes, the QB return the ID of the node
# I use DbAttribute.get_all_values_for_nodepk
# to get the dictionary
return DbAttribute.get_all_values_for_nodepk(res)
elif key == 'extras':
# same as attributes
return DbExtra.get_all_values_for_nodepk(res)
elif key in ('_metadata', 'transport_params'):
# Metadata and transport_params are stored as json strings in the DB:
return json_loads(res)
elif isinstance(res,
(self.Group, self.Node, self.Computer, self.User)):
returnval = res.get_aiida_class()
else:
returnval = res
return returnval
def _db_iterextras(self):
from aiida.backends.djsite.db.models import DbExtra
extraslist = DbExtra.list_all_node_elements(self.dbnode)
for e in extraslist:
yield (e.key, e.getvalue())
def _get_db_extra(self, key, *args):
from aiida.backends.djsite.db.models import DbExtra
return DbExtra.get_value_for_node(dbnode=self.dbnode, key=key)
def _db_store(self, with_transaction=True):
"""
Store a new node in the DB, also saving its repository directory
and attributes.
After being called attributes cannot be
changed anymore! Instead, extras can be changed only AFTER calling
this store() function.
:note: After successful storage, those links that are in the cache, and
for which also the parent node is already stored, will be
automatically stored. The others will remain unstored.
:parameter with_transaction: if False, no transaction is used. This
is meant to be used ONLY if the outer calling function has already
a transaction open!
:param bool use_cache: Whether I attempt to find an equal node in the DB.
"""
# TODO: This needs to be generalized, allowing for flexible methods
# for storing data and its attributes.
from django.db import transaction
from aiida.common.utils import EmptyContextManager
from aiida.common.exceptions import ValidationError
from aiida.backends.djsite.db.models import DbAttribute
import aiida.orm.autogroup
if with_transaction:
context_man = transaction.atomic()
else:
context_man = EmptyContextManager()
# I save the corresponding django entry
# I set the folder
# NOTE: I first store the files, then only if this is successful,
# I store the DB entry. In this way,
# I assume that if a node exists in the DB, its folder is in place.
# On the other hand, periodically the user might need to run some
# bookkeeping utility to check for lone folders.
self._repository_folder.replace_with_folder(
self._get_temp_folder().abspath, move=True, overwrite=True)
# I do the transaction only during storage on DB to avoid timeout
# problems, especially with SQLite
try:
with context_man:
# Save the row
self._dbnode.save()
# Save its attributes 'manually' without incrementing
# the version for each add.
DbAttribute.reset_values_for_node(self._dbnode,
attributes=self._attrs_cache,
with_transaction=False)
# This should not be used anymore: I delete it to
# possibly free memory
del self._attrs_cache
self._temp_folder = None
self._to_be_stored = False
# Here, I store those links that were in the cache and
# that are between stored nodes.
self._store_cached_input_links()
# This is one of the few cases where it is ok to do a 'global'
# except, also because I am re-raising the exception
except:
# I put back the files in the sandbox folder since the
# transaction did not succeed
self._get_temp_folder().replace_with_folder(
self._repository_folder.abspath, move=True, overwrite=True)
raise
from aiida.backends.djsite.db.models import DbExtra
# I store the hash without cleaning and without incrementing the nodeversion number
DbExtra.set_value_for_node(self._dbnode, _HASH_EXTRA_KEY, self.get_hash())
return self
