def del_value(cls, key):
"""Delete a setting value."""
setting = sa.get_scoped_session().query(DbSetting).filter(key=key)
setting.val = None
setting.time = timezone.datetime.utcnow()
flag_modified(setting, 'val')
setting.save()
def test_settings(self):
"""
Test the settings table (similar to Attributes, but without the key.
"""
from aiida.backends.sqlalchemy.models.settings import DbSetting
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
from pytz import UTC
from aiida.utils import timezone
from sqlalchemy.exc import IntegrityError
DbSetting.set_value(key='pippo', value=[1, 2, 3])
# s1 = DbSetting.objects.get(key='pippo')
s1 = DbSetting.query.filter_by(key='pippo').first()
self.assertEqual(s1.getvalue(), [1, 2, 3])
s2 = DbSetting(key='pippo')
s2.time = timezone.datetime.now(tz=UTC)
with self.assertRaises(IntegrityError):
with session.begin_nested():
# same name...
session.add(s2)
# Should replace pippo
DbSetting.set_value(key='pippo', value="a")
s1 = DbSetting.query.filter_by(key='pippo').first()
self.assertEqual(s1.getvalue(), "a")
def get(self, computer, user):
"""
Return a SqlaAuthInfo given a computer and a user
:param computer: a Computer instance
:param user: a User instance
:return: an AuthInfo object associated with the given computer and user
:raise NotExistent: if the user is not configured to use computer
:raise sqlalchemy.orm.exc.MultipleResultsFound: if the user is configured
more than once to use the computer! Should never happen
"""
from aiida.backends.sqlalchemy.models.authinfo import DbAuthInfo
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
from sqlalchemy.orm.exc import MultipleResultsFound, NoResultFound
try:
authinfo = session.query(DbAuthInfo).filter_by(
dbcomputer_id=computer.id,
aiidauser_id=user.id,
).one()
return self.from_dbmodel(authinfo)
except NoResultFound:
raise exceptions.NotExistent(
"The aiida user {} is not configured to use computer {}".format(
user.email, computer.name))
except MultipleResultsFound:
raise exceptions.ConfigurationError(
"The aiida user {} is configured more than once to use "
"computer {}! Only one configuration is allowed".format(
user.email, computer.name))
def set_correct_schema_version_and_backend():
from aiida.utils import timezone
# Setting the correct backend and schema version
SQLA_SCHEMA_VERSION = 0.1
session = sa.get_scoped_session()
with session.begin(subtransactions=True):
# Setting manually the correct schema version
session.execute(
'DELETE FROM db_dbsetting WHERE key=\'db|schemaversion\'')
session.execute(
'INSERT INTO db_dbsetting (key, val, description, time) values '
'(\'db|schemaversion\', \'{}\', '
'\'The version of the schema used in this database.\', \'{}\')'.
format(SQLA_SCHEMA_VERSION, timezone.datetime.now()))
# Setting the correct backend
session.execute('DELETE FROM db_dbsetting WHERE key=\'db|backend\'')
session.execute(
'INSERT INTO db_dbsetting (key, val, description, time) values '
'(\'db|backend\', \'"{}"\', '
'\'The backend used to communicate with database.\', \'{}\')'.
format(BACKEND_SQLA, timezone.datetime.now()))
session.commit()
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
"""
session = get_scoped_session()
if clean:
self.clean_values()
session.add(self._dbmodel)
if links:
for link_triple in links:
self._add_link(*link_triple)
if with_transaction:
try:
session.commit()
except SQLAlchemyError:
session.rollback()
raise
return self
def clear(self):
"""Remove all the nodes from this group."""
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
# Note we have to call `dbmodel` and `_dbmodel` to circumvent the `ModelWrapper`
self.dbmodel.dbnodes = []
session.commit()
def count(self):
"""Return the number of entities in this group.
:return: integer number of entities contained within the group
"""
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
return session.query(self.MODEL_CLASS).join(self.MODEL_CLASS.dbnodes).filter(DbGroup.id == self.pk).count()
def delete(self, commit=True):
"""Emulate the behavior of Django's delete() method
:param commit: whether to do a commit or just remover from the session"""
sess = get_scoped_session()
sess.delete(self)
if commit:
sess.commit()
def _update_db_description_field(self, field_value):
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
self.dbnode.description = field_value
if not self._to_be_stored:
session.add(self._dbnode)
self._increment_version_number_db()
def add_comment(self, content, user=None):
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
if self._to_be_stored:
raise ModificationNotAllowed("Comments can be added only after "
"storing the node")
comment = DbComment(dbnode=self._dbnode, user=user, content=content)
session.add(comment)
try:
session.commit()
except:
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
session.rollback()
raise
def _update_db_label_field(self, field_value):
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
self._dbnode.label = field_value
if self.is_stored:
session.add(self._dbnode)
self._increment_version_number_db()
def add_nodes(self, nodes, **kwargs):
"""Add a node or a set of nodes to the group.
:note: all the nodes *and* the group itself have to be stored.
:param nodes: a list of `BackendNode` instance to be added to this group
:param kwargs:
skip_orm: When the flag is on, the SQLA ORM is skipped and SQLA is used
to create a direct SQL INSERT statement to the group-node relationship
table (to improve speed).
"""
from sqlalchemy.exc import IntegrityError # pylint: disable=import-error, no-name-in-module
from sqlalchemy.dialects.postgresql import insert # pylint: disable=import-error, no-name-in-module
from aiida.orm.implementation.sqlalchemy.nodes import SqlaNode
from aiida.backends.sqlalchemy import get_scoped_session
from aiida.backends.sqlalchemy.models.base import Base
super().add_nodes(nodes)
skip_orm = kwargs.get('skip_orm', False)
def check_node(given_node):
""" Check if given node is of correct type and stored """
if not isinstance(given_node, SqlaNode):
raise TypeError('invalid type {}, has to be {}'.format(type(given_node), SqlaNode))
if not given_node.is_stored:
raise ValueError('At least one of the provided nodes is unstored, stopping...')
with utils.disable_expire_on_commit(get_scoped_session()) as session:
if not skip_orm:
# Get dbnodes here ONCE, otherwise each call to dbnodes will re-read the current value in the database
dbnodes = self._dbmodel.dbnodes
for node in nodes:
check_node(node)
# Use pattern as suggested here:
# http://docs.sqlalchemy.org/en/latest/orm/session_transaction.html#using-savepoint
try:
with session.begin_nested():
dbnodes.append(node.dbmodel)
session.flush()
except IntegrityError:
# Duplicate entry, skip
pass
else:
ins_dict = list()
for node in nodes:
check_node(node)
ins_dict.append({'dbnode_id': node.id, 'dbgroup_id': self.id})
my_table = Base.metadata.tables['db_dbgroup_dbnodes']
ins = insert(my_table).values(ins_dict)
session.execute(ins.on_conflict_do_nothing(index_elements=['dbnode_id', 'dbgroup_id']))
# Commit everything as up till now we've just flushed
session.commit()
def transition_json_column(profile=None):
"""
Migrate the TEXT column containing JSON into JSON columns
"""
print("\nChanging various columns to JSON format.")
if not is_dbenv_loaded():
transition_load_db_env(profile=profile)
table_col = [
('db_dbauthinfo', 'metadata'), ('db_dbauthinfo', 'auth_params'),
('db_dbcomputer', 'metadata'), ('db_dbcomputer', 'transport_params'),
('db_dblog', 'metadata')
]
inspector = reflection.Inspector.from_engine(sa.get_scoped_session().bind)
sql = ("ALTER TABLE {table} ALTER COLUMN {column} TYPE JSONB "
"USING {column}::JSONB")
session = sa.get_scoped_session()
with session.begin(subtransactions=True):
for table, col in table_col:
table_cols = inspector.get_columns(table)
col_type_list = [_["type"] for _ in table_cols if _["name"] == col]
if len(col_type_list) != 1:
raise Exception("Problem with table {} and column {}. Either"
"the column doesn't exist or multiple "
"occurrences were found.".format(table, col))
if isinstance(col_type_list[0], JSON):
print(
"Column {} of table {} is already in JSON format. "
"Proceeding with the next table & column.".format(
col, table))
continue
print("Changing column {} of table {} in JSON format.".format(
table, col))
session.execute(sql.format(table=table, column=col))
session.commit()
def _reset_db_extras(self, new_extras):
try:
self.dbnode.reset_extras(new_extras)
self._increment_version_number_db()
except:
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
session.rollback()
raise
def _del_db_extra(self, key):
try:
self.dbnode.del_extra(key)
self._increment_version_number_db()
except:
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
session.rollback()
raise
def get_dbauthinfo(computer, aiidauser):
"""
Given a computer and a user, returns a DbAuthInfo object
:param computer: a computer (can be a string, Computer or DbComputer instance)
:param aiidauser: a user, can be a DbUser or a User instance
:return: a DbAuthInfo instance
:raise NotExistent: if the user is not configured to use computer
:raise sqlalchemy.orm.exc.MultipleResultsFound: if the user is configured
more than once to use the computer! Should never happen
"""
from aiida.common.exceptions import InternalError
if settings.BACKEND == BACKEND_DJANGO:
from aiida.backends.djsite.db.models import DbComputer, DbAuthInfo, DbUser
from django.core.exceptions import (ObjectDoesNotExist,
MultipleObjectsReturned)
try:
authinfo = DbAuthInfo.objects.get(
# converts from name, Computer or DbComputer instance to
# a DbComputer instance
dbcomputer=DbComputer.get_dbcomputer(computer),
aiidauser=aiidauser.id)
except ObjectDoesNotExist:
raise NotExistent(
"The aiida user {} is not configured to use computer {}".
format(aiidauser.email, computer.name))
except MultipleObjectsReturned:
raise ConfigurationError(
"The aiida user {} is configured more than once to use "
"computer {}! Only one configuration is allowed".format(
aiidauser.email, computer.name))
elif settings.BACKEND == BACKEND_SQLA:
from aiida.backends.sqlalchemy.models.authinfo import DbAuthInfo
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
from sqlalchemy.orm.exc import MultipleResultsFound, NoResultFound
try:
authinfo = session.query(DbAuthInfo).filter_by(
dbcomputer_id=computer.id,
aiidauser_id=aiidauser.id,
).one()
except NoResultFound:
raise NotExistent(
"The aiida user {} is not configured to use computer {}".
format(aiidauser.email, computer.name))
except MultipleResultsFound:
raise ConfigurationError(
"The aiida user {} is configured more than once to use "
"computer {}! Only one configuration is allowed".format(
aiidauser.email, computer.name))
else:
raise InternalError("unknown backend {}".format(settings.BACKEND))
return authinfo
def execute_raw(query):
"""Execute a raw SQL statement and return the result.
:param query: a string containing a raw SQL statement
:return: the result of the query
"""
session = get_scoped_session()
result = session.execute(query)
return result.fetchall()
def validate_table_existence(self):
"""Verify that the `DbSetting` table actually exists.
:raises: `~aiida.common.exceptions.NotExistent` if the settings table does not exist
"""
from sqlalchemy.engine import reflection
inspector = reflection.Inspector.from_engine(get_scoped_session().bind)
if self.table_name not in inspector.get_table_names():
raise NotExistent('the settings table does not exist')
def _increment_version_number_db(self):
self._dbnode.nodeversion = DbNode.nodeversion + 1
try:
self._dbnode.save()
except:
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
session.rollback()
raise
def save(self, commit=True):
"""Emulate the behavior of Django's save() method
:param commit: whether to do a commit or just add to the session
:return: the SQLAlchemy instance"""
sess = get_scoped_session()
sess.add(self)
if commit:
sess.commit()
return self
def table_check_test():
"""
Checks if the db_setting table exists in the database. If it doesn't exist
it rainses a KeyError.
"""
from sqlalchemy.engine import reflection
from aiida.backends import sqlalchemy as sa
inspector = reflection.Inspector.from_engine(get_scoped_session().bind)
if 'db_dbsetting' not in inspector.get_table_names():
raise KeyError("No table found")
def remove(self, authinfo_id):
from sqlalchemy.orm.exc import NoResultFound
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
try:
session.query(DbAuthInfo).filter_by(id=authinfo_id).delete()
session.commit()
except NoResultFound:
raise exceptions.NotExistent("AuthInfo with id '{}' not found".format(authinfo_id))
def delete(self, pk):
try:
session = get_scoped_session()
session.query(DbComputer).get(pk).delete()
session.commit()
except SQLAlchemyError as exc:
raise exceptions.InvalidOperation(
'Unable to delete the requested computer: it is possible that there '
'is at least one node using this computer (original message: {})'
.format(exc))
def remove_nodes(self, nodes):
if not self.is_stored:
raise ModificationNotAllowed("Cannot remove nodes from a group "
"before storing")
from aiida.orm.implementation.sqlalchemy.node import Node
# First convert to a list
if isinstance(nodes, (Node, DbNode)):
nodes = [nodes]
if isinstance(
nodes,
basestring) or not isinstance(nodes, collections.Iterable):
raise TypeError("Invalid type passed as the 'nodes' parameter to "
"remove_nodes, can only be a Node, DbNode, or a "
"list of such objects, it is instead {}".format(
str(type(nodes))))
list_nodes = []
# Have to save dbndoes here ONCE, otherwise it will be reloaded from
# the database each time we access it through _dbgroup.dbnodes
dbnodes = self._dbgroup.dbnodes
for node in nodes:
if not isinstance(node, (Node, DbNode)):
raise TypeError("Invalid type of one of the elements passed "
"to add_nodes, it should be either a Node or "
"a DbNode, it is instead {}".format(
str(type(node))))
if node.id is None:
raise ValueError("At least one of the provided nodes is "
"unstored, stopping...")
if isinstance(node, Node):
node = node.dbnode
# If we don't check first, SqlA might issue a DELETE statement for
# an unexisting key, resulting in an error
if node in dbnodes:
list_nodes.append(node)
for node in list_nodes:
dbnodes.remove(node)
sa.get_scoped_session().commit()
def _remove_comment(self, comment_pk, user):
comment = DbComment.query.filter_by(dbnode=self._dbnode,
id=comment_pk).first()
if comment:
try:
comment.delete()
except:
from aiida.backends.sqlalchemy import get_scoped_session
session = get_scoped_session()
session.rollback()
raise
def remove_nodes(self, nodes, **kwargs):
"""Remove a node or a set of nodes from the group.
:note: all the nodes *and* the group itself have to be stored.
:param nodes: a list of `BackendNode` instance to be added to this group
:param kwargs:
skip_orm: When the flag is set to `True`, the SQLA ORM is skipped and SQLA is used to create a direct SQL
DELETE statement to the group-node relationship table in order to improve speed.
"""
from sqlalchemy import and_
from aiida.backends.sqlalchemy import get_scoped_session
from aiida.backends.sqlalchemy.models.base import Base
from aiida.orm.implementation.sqlalchemy.nodes import SqlaNode
super().remove_nodes(nodes)
# Get dbnodes here ONCE, otherwise each call to dbnodes will re-read the current value in the database
dbnodes = self._dbmodel.dbnodes
skip_orm = kwargs.get('skip_orm', False)
def check_node(node):
if not isinstance(node, SqlaNode):
raise TypeError(
f'invalid type {type(node)}, has to be {SqlaNode}')
if node.id is None:
raise ValueError(
'At least one of the provided nodes is unstored, stopping...'
)
list_nodes = []
with utils.disable_expire_on_commit(get_scoped_session()) as session:
if not skip_orm:
for node in nodes:
check_node(node)
# Check first, if SqlA issues a DELETE statement for an unexisting key it will result in an error
if node.dbmodel in dbnodes:
list_nodes.append(node.dbmodel)
for node in list_nodes:
dbnodes.remove(node)
else:
table = Base.metadata.tables['db_dbgroup_dbnodes']
for node in nodes:
check_node(node)
clause = and_(table.c.dbnode_id == node.id,
table.c.dbgroup_id == self.id)
statement = table.delete().where(clause)
session.execute(statement)
session.commit()
def test_load_nodes(self):
"""
Test for load_node() function.
"""
from aiida.orm import load_node
from aiida.common.exceptions import NotExistent
import aiida.backends.sqlalchemy
from aiida.backends.sqlalchemy import get_scoped_session
a = Node()
a.store()
self.assertEquals(a.pk, load_node(node_id=a.pk).pk)
self.assertEquals(a.pk, load_node(node_id=a.uuid).pk)
self.assertEquals(a.pk, load_node(pk=a.pk).pk)
self.assertEquals(a.pk, load_node(uuid=a.uuid).pk)
session = get_scoped_session()
try:
session.begin_nested()
with self.assertRaises(ValueError):
load_node(node_id=a.pk, pk=a.pk)
finally:
session.rollback()
try:
session.begin_nested()
with self.assertRaises(ValueError):
load_node(pk=a.pk, uuid=a.uuid)
finally:
session.rollback()
try:
session.begin_nested()
with self.assertRaises(ValueError):
load_node(pk=a.uuid)
finally:
session.rollback()
try:
session.begin_nested()
with self.assertRaises(ValueError):
load_node(uuid=a.pk)
finally:
session.rollback()
try:
session.begin_nested()
with self.assertRaises(ValueError):
load_node()
finally:
session.rollback()
def delete(self):
session = sa.get_scoped_session()
if self.pk is not None:
session.delete(self._dbgroup)
session.commit()
new_group = copy(self._dbgroup)
make_transient(new_group)
new_group.id = None
self._dbgroup = new_group
def setUp(self):
"""Go to a specific schema version before running tests."""
from aiida.backends import sqlalchemy as sa
from aiida.orm import autogroup
self.current_autogroup = autogroup.current_autogroup
autogroup.current_autogroup = None
assert self.migrate_from and self.migrate_to, \
"TestCase '{}' must define migrate_from and migrate_to properties".format(type(self).__name__)
self.migrate_from = [(self.app, self.migrate_from)]
self.migrate_to = [(self.app, self.migrate_to)]
executor = MigrationExecutor(connection)
self.apps = executor.loader.project_state(self.migrate_from).apps
self.schema_editor = connection.schema_editor()
# Reset session for the migration
sa.get_scoped_session().close()
# Reverse to the original migration
with Capturing():
executor.migrate(self.migrate_from)
# Reset session after the migration
sa.get_scoped_session().close()
self.DbLink = self.apps.get_model('db', 'DbLink')
self.DbNode = self.apps.get_model('db', 'DbNode')
self.DbUser = self.apps.get_model('db', 'DbUser')
self.DbUser.objects.all().delete()
self.default_user = self.DbUser(1, 'aiida@localhost')
self.default_user.save()
try:
self.setUpBeforeMigration()
# Run the migration to test
executor = MigrationExecutor(connection)
executor.loader.build_graph()
# Reset session for the migration
sa.get_scoped_session().close()
with Capturing():
executor.migrate(self.migrate_to)
# Reset session after the migration
sa.get_scoped_session().close()
self.apps = executor.loader.project_state(self.migrate_to).apps
except Exception:
# Bring back the DB to the correct state if this setup part fails
import traceback
traceback.print_stack()
self._revert_database_schema()
raise
def add_nodes(self, nodes):
from sqlalchemy.exc import IntegrityError
if not self.is_stored:
raise ModificationNotAllowed("Cannot add nodes to a group before "
"storing")
from aiida.orm.implementation.sqlalchemy.node import Node
from aiida.backends.sqlalchemy import get_scoped_session
# First convert to a list
if isinstance(nodes, (Node, DbNode)):
nodes = [nodes]
if isinstance(
nodes,
basestring) or not isinstance(nodes, collections.Iterable):
raise TypeError("Invalid type passed as the 'nodes' parameter to "
"add_nodes, can only be a Node, DbNode, or a list "
"of such objects, it is instead {}".format(
str(type(nodes))))
with utils.disable_expire_on_commit(get_scoped_session()) as session:
# Get dbnodes here ONCE, otherwise each call to _dbgroup.dbnodes will
# re-read the current value in the database
dbnodes = self._dbgroup.dbnodes
for node in nodes:
if not isinstance(node, (Node, DbNode)):
raise TypeError(
"Invalid type of one of the elements passed "
"to add_nodes, it should be either a Node or "
"a DbNode, it is instead {}".format(str(type(node))))
if node.id is None:
raise ValueError("At least one of the provided nodes is "
"unstored, stopping...")
if isinstance(node, Node):
to_add = node.dbnode
else:
to_add = node
# Use pattern as suggested here:
# http://docs.sqlalchemy.org/en/latest/orm/session_transaction.html#using-savepoint
try:
with session.begin_nested():
dbnodes.append(to_add)
session.flush()
except IntegrityError:
# Duplicate entry, skip
pass
# Commit everything as up till now we've just flushed
session.commit()
