# -*- coding: utf-8 -*- ########################################################################### # Copyright (c), The AiiDA team. All rights reserved. # # This file is part of the AiiDA code. # # # # The code is hosted on GitHub at https://github.com/aiidateam/aiida-core # # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### """Django implementation of the `BackendNode` and `BackendNodeCollection` classes.""" # pylint: disable=import-error,no-name-in-module from datetime import datetime from django.core.exceptions import ObjectDoesNotExist from django.db import transaction, IntegrityError from aiida.backends.djsite.db import models from aiida.common import exceptions from aiida.common.lang import type_check from aiida.orm.utils.node import clean_value from .. import BackendNode, BackendNodeCollection from . import entities from . import utils as dj_utils from .computers import DjangoComputer from .users import DjangoUser [docs]class DjangoNode(entities.DjangoModelEntity[models.DbNode], BackendNode): """Django Node backend entity""" # pylint: disable=too-many-public-methods MODEL_CLASS = models.DbNode LINK_CLASS = models.DbLink [docs] def __init__( self, backend, node_type, user, computer=None, process_type=None, label='', description='', ctime=None, mtime=None ): """Construct a new `BackendNode` instance wrapping a new `DbNode` instance. :param backend: the backend :param node_type: the node type string :param user: associated `BackendUser` :param computer: associated `BackendComputer` :param label: string label :param description: string description :param ctime: The creation time as datetime object :param mtime: The modification time as datetime object """ # pylint: disable=too-many-arguments super().__init__(backend) arguments = { 'user': user.dbmodel, 'node_type': node_type, 'process_type': process_type, 'label': label, 'description': description, } type_check(user, DjangoUser) if computer: type_check(computer, DjangoComputer, 'computer is of type {}'.format(type(computer))) arguments['dbcomputer'] = computer.dbmodel if ctime: type_check(ctime, datetime, 'the given ctime is of type {}'.format(type(ctime))) arguments['ctime'] = ctime if mtime: type_check(mtime, datetime, 'the given mtime is of type {}'.format(type(mtime))) arguments['mtime'] = mtime self._dbmodel = dj_utils.ModelWrapper(models.DbNode(**arguments)) [docs] def clone(self): """Return an unstored clone of ourselves. :return: an unstored `BackendNode` with the exact same attributes and extras as self """ arguments = { 'node_type': self._dbmodel.node_type, 'process_type': self._dbmodel.process_type, 'user': self._dbmodel.user, 'dbcomputer': self._dbmodel.dbcomputer, 'label': self._dbmodel.label, 'description': self._dbmodel.description, } clone = self.__class__.__new__(self.__class__) # pylint: disable=no-value-for-parameter clone.__init__(self.backend, self.node_type, self.user) clone._dbmodel = dj_utils.ModelWrapper(models.DbNode(**arguments)) # pylint: disable=protected-access return clone @property def computer(self): """Return the computer of this node. :return: the computer or None :rtype: `BackendComputer` or None """ try: return self.backend.computers.from_dbmodel(self._dbmodel.dbcomputer) except TypeError: return None @computer.setter def computer(self, computer): """Set the computer of this node. :param computer: a `BackendComputer` """ type_check(computer, DjangoComputer, allow_none=True) if computer is not None: computer = computer.dbmodel self._dbmodel.dbcomputer = computer @property def user(self): """Return the user of this node. :return: the user :rtype: `BackendUser` """ return self.backend.users.from_dbmodel(self._dbmodel.user) @user.setter def user(self, user): """Set the user of this node. :param user: a `BackendUser` """ type_check(user, DjangoUser) self._dbmodel.user = user.dbmodel @property def attributes(self): """Return the complete attributes dictionary. .. warning:: While the node is unstored, this will return references of the attributes on the database model, meaning that changes on the returned values (if they are mutable themselves, e.g. a list or dictionary) will automatically be reflected on the database model as well. As soon as the node is stored, the returned attributes will be a deep copy and mutations of the database attributes will have to go through the appropriate set methods. Therefore, once stored, retrieving a deep copy can be a heavy operation. If you only need the keys or some values, use the iterators `attributes_keys` and `attributes_items`, or the getters `get_attribute` and `get_attribute_many` instead. :return: the attributes as a dictionary """ return self.dbmodel.attributes [docs] def get_attribute(self, key): """Return the value of an attribute. .. warning:: While the node is unstored, this will return a reference of the attribute on the database model, meaning that changes on the returned value (if they are mutable themselves, e.g. a list or dictionary) will automatically be reflected on the database model as well. As soon as the node is stored, the returned attribute will be a deep copy and mutations of the database attributes will have to go through the appropriate set methods. :param key: name of the attribute :return: the value of the attribute :raises AttributeError: if the attribute does not exist """ try: return self._dbmodel.attributes[key] except KeyError as exception: raise AttributeError('attribute `{}` does not exist'.format(exception)) [docs] def get_attribute_many(self, keys): """Return the values of multiple attributes. .. warning:: While the node is unstored, this will return references of the attributes on the database model, meaning that changes on the returned values (if they are mutable themselves, e.g. a list or dictionary) will automatically be reflected on the database model as well. As soon as the node is stored, the returned attributes will be a deep copy and mutations of the database attributes will have to go through the appropriate set methods. Therefore, once stored, retrieving a deep copy can be a heavy operation. If you only need the keys or some values, use the iterators `attributes_keys` and `attributes_items`, or the getters `get_attribute` and `get_attribute_many` instead. :param keys: a list of attribute names :return: a list of attribute values :raises AttributeError: if at least one attribute does not exist """ try: return [self.get_attribute(key) for key in keys] except KeyError as exception: raise AttributeError('attribute `{}` does not exist'.format(exception)) [docs] def set_attribute(self, key, value): """Set an attribute to the given value. :param key: name of the attribute :param value: value of the attribute """ if self.is_stored: value = clean_value(value) self._dbmodel.attributes[key] = value self._flush_if_stored({'attributes'}) [docs] def set_attribute_many(self, attributes): """Set multiple attributes. .. note:: This will override any existing attributes that are present in the new dictionary. :param attributes: a dictionary with the attributes to set """ if self.is_stored: attributes = {key: clean_value(value) for key, value in attributes.items()} for key, value in attributes.items(): # We need to use `self.dbmodel` without the underscore, because otherwise the second iteration will refetch # what is in the database and we lose the initial changes. self.dbmodel.attributes[key] = value self._flush_if_stored({'attributes'}) [docs] def reset_attributes(self, attributes): """Reset the attributes. .. note:: This will completely clear any existing attributes and replace them with the new dictionary. :param attributes: a dictionary with the attributes to set """ if self.is_stored: attributes = clean_value(attributes) self.dbmodel.attributes = attributes self._flush_if_stored({'attributes'}) [docs] def delete_attribute(self, key): """Delete an attribute. :param key: name of the attribute :raises AttributeError: if the attribute does not exist """ try: self._dbmodel.attributes.pop(key) except KeyError as exception: raise AttributeError('attribute `{}` does not exist'.format(exception)) else: self._flush_if_stored({'attributes'}) [docs] def delete_attribute_many(self, keys): """Delete multiple attributes. :param keys: names of the attributes to delete :raises AttributeError: if at least one of the attribute does not exist """ non_existing_keys = [key for key in keys if key not in self._dbmodel.attributes] if non_existing_keys: raise AttributeError('attributes `{}` do not exist'.format(', '.join(non_existing_keys))) for key in keys: self.dbmodel.attributes.pop(key) self._flush_if_stored({'attributes'}) [docs] def clear_attributes(self): """Delete all attributes.""" self._dbmodel.attributes = {} self._flush_if_stored({'attributes'}) [docs] def attributes_items(self): """Return an iterator over the attributes. :return: an iterator with attribute key value pairs """ for key, value in self._dbmodel.attributes.items(): yield key, value [docs] def attributes_keys(self): """Return an iterator over the attribute keys. :return: an iterator with attribute keys """ for key in self._dbmodel.attributes: yield key @property def extras(self): """Return the complete extras dictionary. .. warning:: While the node is unstored, this will return references of the extras on the database model, meaning that changes on the returned values (if they are mutable themselves, e.g. a list or dictionary) will automatically be reflected on the database model as well. As soon as the node is stored, the returned extras will be a deep copy and mutations of the database extras will have to go through the appropriate set methods. Therefore, once stored, retrieving a deep copy can be a heavy operation. If you only need the keys or some values, use the iterators `extras_keys` and `extras_items`, or the getters `get_extra` and `get_extra_many` instead. :return: the extras as a dictionary """ return self.dbmodel.extras [docs] def get_extra(self, key): """Return the value of an extra. .. warning:: While the node is unstored, this will return a reference of the extra on the database model, meaning that changes on the returned value (if they are mutable themselves, e.g. a list or dictionary) will automatically be reflected on the database model as well. As soon as the node is stored, the returned extra will be a deep copy and mutations of the database extras will have to go through the appropriate set methods. :param key: name of the extra :return: the value of the extra :raises AttributeError: if the extra does not exist """ try: return self._dbmodel.extras[key] except KeyError as exception: raise AttributeError('extra `{}` does not exist'.format(exception)) [docs] def get_extra_many(self, keys): """Return the values of multiple extras. .. warning:: While the node is unstored, this will return references of the extras on the database model, meaning that changes on the returned values (if they are mutable themselves, e.g. a list or dictionary) will automatically be reflected on the database model as well. As soon as the node is stored, the returned extras will be a deep copy and mutations of the database extras will have to go through the appropriate set methods. Therefore, once stored, retrieving a deep copy can be a heavy operation. If you only need the keys or some values, use the iterators `extras_keys` and `extras_items`, or the getters `get_extra` and `get_extra_many` instead. :param keys: a list of extra names :return: a list of extra values :raises AttributeError: if at least one extra does not exist """ try: return [self.get_extra(key) for key in keys] except KeyError as exception: raise AttributeError('extra `{}` does not exist'.format(exception)) [docs] def set_extra(self, key, value): """Set an extra to the given value. :param key: name of the extra :param value: value of the extra """ if self.is_stored: value = clean_value(value) self._dbmodel.extras[key] = value self._flush_if_stored({'extras'}) [docs] def set_extra_many(self, extras): """Set multiple extras. .. note:: This will override any existing extras that are present in the new dictionary. :param extras: a dictionary with the extras to set """ if self.is_stored: extras = {key: clean_value(value) for key, value in extras.items()} for key, value in extras.items(): self.dbmodel.extras[key] = value self._flush_if_stored({'extras'}) [docs] def reset_extras(self, extras): """Reset the extras. .. note:: This will completely clear any existing extras and replace them with the new dictionary. :param extras: a dictionary with the extras to set """ if self.is_stored: extras = clean_value(extras) self.dbmodel.extras = extras self._flush_if_stored({'extras'}) [docs] def delete_extra(self, key): """Delete an extra. :param key: name of the extra :raises AttributeError: if the extra does not exist """ try: self._dbmodel.extras.pop(key) except KeyError as exception: raise AttributeError('extra `{}` does not exist'.format(exception)) else: self._flush_if_stored({'extras'}) [docs] def delete_extra_many(self, keys): """Delete multiple extras. :param keys: names of the extras to delete :raises AttributeError: if at least one of the extra does not exist """ non_existing_keys = [key for key in keys if key not in self._dbmodel.extras] if non_existing_keys: raise AttributeError('extras `{}` do not exist'.format(', '.join(non_existing_keys))) for key in keys: self.dbmodel.extras.pop(key) self._flush_if_stored({'extras'}) [docs] def clear_extras(self): """Delete all extras.""" self._dbmodel.extras = {} self._flush_if_stored({'extras'}) [docs] def extras_items(self): """Return an iterator over the extras. :return: an iterator with extra key value pairs """ for key, value in self._dbmodel.extras.items(): yield key, value [docs] def extras_keys(self): """Return an iterator over the extra keys. :return: an iterator with extra keys """ for key in self._dbmodel.extras: yield key [docs] def _flush_if_stored(self, fields=None): if self._dbmodel.is_saved(): self._dbmodel._flush(fields) # pylint: disable=protected-access [docs] def add_incoming(self, source, link_type, link_label): """Add a link of the given type from a given node to ourself. :param source: the node from which the link is coming :param link_type: the link type :param link_label: the link label :return: True if the proposed link is allowed, False otherwise :raise aiida.common.ModificationNotAllowed: if either source or target node is not stored """ type_check(source, DjangoNode) if not self.is_stored: raise exceptions.ModificationNotAllowed('node has to be stored when adding an incoming link') if not source.is_stored: raise exceptions.ModificationNotAllowed('source node has to be stored when adding a link from it') self._add_link(source, link_type, link_label) [docs] def _add_link(self, source, link_type, link_label): """Add a link of the given type from a given node to ourself. :param source: the node from which the link is coming :param link_type: the link type :param link_label: the link label """ savepoint_id = None try: # Transactions are needed here for Postgresql: # https://docs.djangoproject.com/en/1.5/topics/db/transactions/#handling-exceptions-within-postgresql-transactions savepoint_id = transaction.savepoint() self.LINK_CLASS(input_id=source.id, output_id=self.id, label=link_label, type=link_type.value).save() transaction.savepoint_commit(savepoint_id) except IntegrityError as exception: transaction.savepoint_rollback(savepoint_id) raise exceptions.UniquenessError('failed to create the link: {}'.format(exception)) [docs] def clean_values(self): self._dbmodel.attributes = clean_value(self._dbmodel.attributes) self._dbmodel.extras = clean_value(self._dbmodel.extras) [docs] def store(self, links=None, with_transaction=True, clean=True): # pylint: disable=arguments-differ """Store the node in the database. :param links: optional links to add before storing :param with_transaction: if False, do not use a transaction because the caller will already have opened one. :param clean: boolean, if True, will clean the attributes and extras before attempting to store """ import contextlib from aiida.backends.djsite.db.models import suppress_auto_now if clean: self.clean_values() # `contextlib.suppress` provides empty context and can be replaced with `contextlib.nullcontext` after we drop # support for python 3.6 with transaction.atomic() if with_transaction else contextlib.suppress(): with suppress_auto_now([(models.DbNode, ['mtime'])]) if self.mtime else contextlib.suppress(): # We need to save the node model instance itself first such that it has a pk # that can be used in the foreign keys that will be needed for setting the # attributes and links self.dbmodel.save() if links: for link_triple in links: self._add_link(*link_triple) return self [docs]class DjangoNodeCollection(BackendNodeCollection): """The collection of Node entries.""" ENTITY_CLASS = DjangoNode [docs] def get(self, pk): """Return a Node entry from the collection with the given id :param pk: id of the node """ try: return self.ENTITY_CLASS.from_dbmodel(models.DbNode.objects.get(pk=pk), self.backend) except ObjectDoesNotExist: raise exceptions.NotExistent("Node with pk '{}' not found".format(pk)) [docs] def delete(self, pk): """Remove a Node entry from the collection with the given id :param pk: id of the node to delete """ try: models.DbNode.objects.filter(pk=pk).delete() # pylint: disable=no-member except ObjectDoesNotExist: raise exceptions.NotExistent("Node with pk '{}' not found".format(pk))