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)
def execute_alembic_command(command_name, **kwargs):
"""Execute an Alembic CLI command.
:param command_name: the sub command name
:param kwargs: parameters to pass to the command
"""
from aiida.backends.sqlalchemy.manager import SqlaBackendManager
manager = SqlaBackendManager()
with manager.alembic_config() as config:
command = getattr(alembic.command, command_name)
command(config, **kwargs)
def get_backend_manager(backend):
"""Get an instance of the `BackendManager` for the current backend.
:param backend: the type of the database backend
:return: `BackendManager`
"""
if backend == BACKEND_DJANGO:
from aiida.backends.djsite.manager import DjangoBackendManager
return DjangoBackendManager()
if backend == BACKEND_SQLA:
from aiida.backends.sqlalchemy.manager import SqlaBackendManager
return SqlaBackendManager()
raise Exception(f'unknown backend type `{backend}`')
class SqlaBackend(SqlBackend[base.Base]):
"""SqlAlchemy implementation of `aiida.orm.implementation.backends.Backend`."""
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)
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
def query(self):
return querybuilder.SqlaQueryBuilder(self)
@property
def users(self):
return self._users
@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()
@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`
def get_backend_entity(self, model):
"""Return a `BackendEntity` instance from a `DbModel` instance."""
return convert.get_backend_entity(model, self)
@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()
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
@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()