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()