# -*- 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 #
###########################################################################
"""SqlAlchemy implementation of `aiida.orm.implementation.backends.Backend`."""
from contextlib import contextmanager
from aiida.backends.sqlalchemy.models import base
from aiida.backends.sqlalchemy.queries import SqlaQueryManager
from aiida.backends.sqlalchemy.manager import SqlaBackendManager
from ..sql import SqlBackend
from . import authinfos
from . import comments
from . import computers
from . import convert
from . import groups
from . import logs
from . import nodes
from . import querybuilder
from . import users
__all__ = ('SqlaBackend',)
[docs]class SqlaBackend(SqlBackend[base.Base]):
"""SqlAlchemy implementation of `aiida.orm.implementation.backends.Backend`."""
[docs] def __init__(self):
"""Construct the backend instance by initializing all the collections."""
self._authinfos = authinfos.SqlaAuthInfoCollection(self)
self._comments = comments.SqlaCommentCollection(self)
self._computers = computers.SqlaComputerCollection(self)
self._groups = groups.SqlaGroupCollection(self)
self._logs = logs.SqlaLogCollection(self)
self._nodes = nodes.SqlaNodeCollection(self)
self._query_manager = SqlaQueryManager(self)
self._schema_manager = SqlaBackendManager()
self._users = users.SqlaUserCollection(self)
[docs] def migrate(self):
self._schema_manager.migrate()
@property
def authinfos(self):
return self._authinfos
@property
def comments(self):
return self._comments
@property
def computers(self):
return self._computers
@property
def groups(self):
return self._groups
@property
def logs(self):
return self._logs
@property
def nodes(self):
return self._nodes
@property
def query_manager(self):
return self._query_manager
[docs] def query(self):
return querybuilder.SqlaQueryBuilder(self)
@property
def users(self):
return self._users
[docs] @contextmanager
def transaction(self):
"""Open a transaction to be used as a context manager.
If there is an exception within the context then the changes will be rolled back and the state will be as before
entering. Transactions can be nested.
"""
session = self.get_session()
nested = session.transaction.nested
try:
session.begin_nested()
yield session
session.commit()
except Exception:
session.rollback()
raise
finally:
if not nested:
# Make sure to commit the outermost session
session.commit()
[docs] @staticmethod
def get_session():
"""Return a database session that can be used by the `QueryBuilder` to perform its query.
:return: an instance of :class:`sqlalchemy.orm.session.Session`
"""
from aiida.backends.sqlalchemy import get_scoped_session
return get_scoped_session()
# Below are abstract methods inherited from `aiida.orm.implementation.sql.backends.SqlBackend`
[docs] def get_backend_entity(self, model):
"""Return a `BackendEntity` instance from a `DbModel` instance."""
return convert.get_backend_entity(model, self)
[docs] @contextmanager
def cursor(self):
"""Return a psycopg cursor to be used in a context manager.
:return: a psycopg cursor
:rtype: :class:`psycopg2.extensions.cursor`
"""
from aiida.backends import sqlalchemy as sa
try:
connection = sa.ENGINE.raw_connection()
yield connection.cursor()
finally:
self.get_connection().close()
[docs] def execute_raw(self, 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
"""
from sqlalchemy.exc import ResourceClosedError # pylint: disable=import-error,no-name-in-module
with self.transaction() as session:
queryset = session.execute(query)
try:
results = queryset.fetchall()
except ResourceClosedError:
return None
return results
[docs] @staticmethod
def get_connection():
"""Get the SQLA database connection
:return: the SQLA database connection
"""
from aiida.backends import sqlalchemy as sa
return sa.ENGINE.raw_connection()