Source code for aiida.storage.psql_dos.orm.entities
###########################################################################
# 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 #
###########################################################################
"""Classes and methods for Django specific backend entities"""
from typing import Generic, Set, TypeVar
from aiida.common.lang import type_check
from aiida.storage.psql_dos.models.base import Base
from . import utils
ModelType = TypeVar('ModelType')
SelfType = TypeVar('SelfType', bound='SqlaModelEntity')
[docs]
class SqlaModelEntity(Generic[ModelType]):
"""A mixin that adds some common SQLA backend entity methods"""
MODEL_CLASS = None
_model: utils.ModelWrapper
[docs]
@classmethod
def _class_check(cls):
"""Assert that the class is correctly configured"""
assert issubclass(cls.MODEL_CLASS, Base), 'Must set the MODEL_CLASS in the derived class to a SQLA model'
[docs]
@classmethod
def from_dbmodel(cls, dbmodel, backend):
"""Create an AiiDA Entity from the corresponding SQLA ORM model and storage backend
:param dbmodel: the SQLAlchemy model to create the entity from
:param backend: the corresponding storage backend
:return: the AiiDA entity
"""
from ..backend import PsqlDosBackend
cls._class_check()
type_check(dbmodel, cls.MODEL_CLASS)
type_check(backend, PsqlDosBackend)
entity = cls.__new__(cls)
super(SqlaModelEntity, entity).__init__(backend)
entity._model = utils.ModelWrapper(dbmodel, backend)
return entity
[docs]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._class_check()
@property
def model(self) -> utils.ModelWrapper:
"""Return an ORM model that correctly updates and flushes the data when getting or setting a field."""
return self._model
@property
def bare_model(self):
"""Return the underlying SQLA ORM model for this entity.
.. warning:: Getting/setting attributes on this model bypasses AiiDA's internal update/flush mechanisms.
"""
return self.model._model
@property
def id(self) -> int:
"""Get the id of this entity
:return: the entity id
"""
return self.model.id
@property
def is_stored(self) -> bool:
"""Is this entity stored?
:return: True if stored, False otherwise
"""
return self.model.id is not None
[docs]
def store(self: SelfType) -> SelfType:
"""Store this entity
:return: the entity itself
"""
self.model.save()
return self
[docs]
def _flush_if_stored(self, fields: Set[str]) -> None:
if self.model.is_saved():
self.model._flush(fields)