diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5321bd6b0b..ff4ed4e025 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -36,12 +36,9 @@ aiida/backends/djsite/db/migrations/__init__.py| aiida/backends/djsite/db/models.py| aiida/backends/djsite/db/subtests/test_generic.py| - aiida/backends/djsite/__init__.py| - aiida/backends/djsite/manage.py| aiida/backends/djsite/queries.py| aiida/backends/profile.py| aiida/backends/general/abstractqueries.py| - aiida/backends/sqlalchemy/__init__.py| aiida/backends/sqlalchemy/migrations/env.py| aiida/backends/sqlalchemy/migrations/versions/0aebbeab274d_base_data_plugin_type_string.py| aiida/backends/sqlalchemy/migrations/versions/35d4ee9a1b0e_code_hidden_attr_to_extra.py| @@ -75,7 +72,6 @@ aiida/backends/sqlalchemy/tests/test_session.py| aiida/backends/sqlalchemy/tests/testbase.py| aiida/backends/sqlalchemy/tests/test_utils.py| - aiida/backends/sqlalchemy/utils.py| aiida/backends/testbase.py| aiida/backends/testimplbase.py| aiida/backends/tests/test_base_dataclasses.py| diff --git a/aiida/backends/__init__.py b/aiida/backends/__init__.py index 606e3bd7c2..a4d87dc829 100644 --- a/aiida/backends/__init__.py +++ b/aiida/backends/__init__.py @@ -14,3 +14,20 @@ BACKEND_DJANGO = 'django' BACKEND_SQLA = 'sqlalchemy' + + +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('unknown backend type `{}`'.format(backend)) diff --git a/aiida/backends/djsite/db/migrations/__init__.py b/aiida/backends/djsite/db/migrations/__init__.py index fe8cf402d5..7c71f7bbfd 100644 --- a/aiida/backends/djsite/db/migrations/__init__.py +++ b/aiida/backends/djsite/db/migrations/__init__.py @@ -16,7 +16,8 @@ from django.core.exceptions import ObjectDoesNotExist from aiida.common.exceptions import AiidaException, DbContentError from six.moves import range -from aiida.backends.djsite.utils import SCHEMA_VERSION_DB_KEY, SCHEMA_VERSION_DB_DESCRIPTION + +from aiida.backends.manager import SCHEMA_VERSION_KEY, SCHEMA_VERSION_DESCRIPTION, SCHEMA_GENERATION_KEY, SCHEMA_GENERATION_DESCRIPTION class DeserializationException(AiidaException): @@ -33,22 +34,39 @@ def _update_schema_version(version, apps, schema_editor): settings table to JSONB) """ db_setting_model = apps.get_model('db', 'DbSetting') - res = db_setting_model.objects.filter(key=SCHEMA_VERSION_DB_KEY).first() + result = db_setting_model.objects.filter(key=SCHEMA_VERSION_KEY).first() # If there is no schema record, create ones - if res is None: - res = db_setting_model() - res.key = SCHEMA_VERSION_DB_KEY - res.description = SCHEMA_VERSION_DB_DESCRIPTION + if result is None: + result = db_setting_model() + result.key = SCHEMA_VERSION_KEY + result.description = SCHEMA_VERSION_DESCRIPTION # If it stores the values in an EAV format, add the value in the tval field - if hasattr(res, 'tval'): - res.tval = str(version) + if hasattr(result, 'tval'): + result.tval = str(version) # Otherwise add it to the val (JSON) fiels else: - res.val = str(version) + result.val = str(version) + + result.save() + + +def _upgrade_schema_generation(version, apps, schema_editor): + """ + The update schema uses the current models (and checks if the value is stored in EAV mode or JSONB) + to avoid to use the DbSettings schema that may change (as it changed with the migration of the + settings table to JSONB) + """ + db_setting_model = apps.get_model('db', 'DbSetting') + result = db_setting_model.objects.filter(key=SCHEMA_GENERATION_KEY).first() + # If there is no schema record, create ones + if result is None: + result = db_setting_model() + result.key = SCHEMA_GENERATION_KEY + result.description = SCHEMA_GENERATION_DESCRIPTION - # Store the final result - res.save() + result.val = str(version) + result.save() def upgrade_schema_version(up_revision, down_revision): @@ -385,8 +403,18 @@ def validate_key(self, key): :return: None if the key is valid :raise aiida.common.ValidationError: if the key is not valid """ - from aiida.backends.utils import validate_attribute_key - return validate_attribute_key(key) + from aiida.backends.utils import AIIDA_ATTRIBUTE_SEP + from aiida.common.exceptions import ValidationError + + if not isinstance(key, six.string_types): + raise ValidationError('The key must be a string.') + if not key: + raise ValidationError('The key cannot be an empty string.') + if AIIDA_ATTRIBUTE_SEP in key: + raise ValidationError( + "The separator symbol '{}' cannot be present " + 'in the key of attributes, extras, etc.'.format(AIIDA_ATTRIBUTE_SEP) + ) def get_value_for_node(self, dbnode, key): """ diff --git a/aiida/backends/djsite/db/subtests/migrations/test_migrations_0037_attributes_extras_settings_json.py b/aiida/backends/djsite/db/subtests/migrations/test_migrations_0037_attributes_extras_settings_json.py index 89ac88d159..a6685144ea 100644 --- a/aiida/backends/djsite/db/subtests/migrations/test_migrations_0037_attributes_extras_settings_json.py +++ b/aiida/backends/djsite/db/subtests/migrations/test_migrations_0037_attributes_extras_settings_json.py @@ -292,8 +292,18 @@ def validate_key(cls, key): :return: None if the key is valid :raise aiida.common.ValidationError: if the key is not valid """ - from aiida.backends.utils import validate_attribute_key - return validate_attribute_key(key) + from aiida.backends.utils import AIIDA_ATTRIBUTE_SEP + from aiida.common.exceptions import ValidationError + + if not isinstance(key, six.string_types): + raise ValidationError('The key must be a string.') + if not key: + raise ValidationError('The key cannot be an empty string.') + if AIIDA_ATTRIBUTE_SEP in key: + raise ValidationError( + "The separator symbol '{}' cannot be present " + 'in the key of attributes, extras, etc.'.format(AIIDA_ATTRIBUTE_SEP) + ) @classmethod def set_value( diff --git a/aiida/backends/djsite/manage.py b/aiida/backends/djsite/manage.py index b4c7676666..4489e9def9 100755 --- a/aiida/backends/djsite/manage.py +++ b/aiida/backends/djsite/manage.py @@ -8,29 +8,27 @@ # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### +# pylint: disable=import-error,no-name-in-module,unused-argument +"""Simple wrapper around Django's `manage.py` CLI script.""" from __future__ import division from __future__ import print_function from __future__ import absolute_import import click -from aiida.cmdline.params import options +from aiida.cmdline.params import options, types @click.command() -@options.PROFILE(required=True) +@options.PROFILE(required=True, type=types.ProfileParamType(load_profile=True)) @click.argument('command', nargs=-1) def main(profile, command): """Simple wrapper around the Django command line tool that first loads an AiiDA profile.""" from django.core.management import execute_from_command_line - - # Load the general load_dbenv. - from aiida.manage.configuration import load_profile from aiida.manage.manager import get_manager - load_profile(profile=profile.name) manager = get_manager() - manager._load_backend(schema_check=False) + manager._load_backend(schema_check=False) # pylint: disable=protected-access # The `execute_from_command` expects a list of command line arguments where the first is the program name that one # would normally call directly. Since this is now replaced by our `click` command we just spoof a random name. @@ -39,4 +37,4 @@ def main(profile, command): if __name__ == '__main__': - main() + main() # pylint: disable=no-value-for-parameter diff --git a/aiida/backends/djsite/manager.py b/aiida/backends/djsite/manager.py new file mode 100644 index 0000000000..b5a9460db9 --- /dev/null +++ b/aiida/backends/djsite/manager.py @@ -0,0 +1,190 @@ +# -*- coding: utf-8 -*- +# pylint: disable=import-error,no-name-in-module +"""Utilities and configuration of the Django database schema.""" +from __future__ import absolute_import + +import os +import django + +from aiida.common import NotExistent +from ..manager import BackendManager, SettingsManager, Setting, SCHEMA_VERSION_KEY, SCHEMA_VERSION_DESCRIPTION + +# The database schema version required to perform schema reset for a given code schema generation +SCHEMA_VERSION_RESET = {'1': None} + + +class DjangoBackendManager(BackendManager): + """Class to manage the database schema.""" + + def get_settings_manager(self): + """Return an instance of the `SettingsManager`. + + :return: `SettingsManager` + """ + if self._settings_manager is None: + self._settings_manager = DjangoSettingsManager() + + return self._settings_manager + + def _load_backend_environment(self): + """Load the backend environment.""" + os.environ['DJANGO_SETTINGS_MODULE'] = 'aiida.backends.djsite.settings' + django.setup() # pylint: disable=no-member + + def reset_backend_environment(self): + """Reset the backend environment.""" + + def is_database_schema_ahead(self): + """Determine whether the database schema version is ahead of the code schema version. + + .. warning:: this will not check whether the schema generations are equal + + :return: boolean, True if the database schema version is ahead of the code schema version. + """ + # For Django the versions numbers are numerical so we can compare them + from distutils.version import StrictVersion + return StrictVersion(self.get_schema_version_database()) > StrictVersion(self.get_schema_version_code()) + + def get_schema_version_code(self): + """Return the code schema version.""" + from .db.models import SCHEMA_VERSION + return SCHEMA_VERSION + + def get_schema_version_reset(self, schema_generation_code): + """Return schema version the database should have to be able to automatically reset to code schema generation. + + :param schema_generation_code: the schema generation of the code. + :return: schema version + """ + return SCHEMA_VERSION_RESET[schema_generation_code] + + def get_schema_generation_database(self): + """Return the database schema version. + + :return: `distutils.version.StrictVersion` with schema version of the database + """ + from django.db.utils import ProgrammingError + from aiida.manage.manager import get_manager + + backend = get_manager()._load_backend(schema_check=False) # pylint: disable=protected-access + + try: + result = backend.execute_raw(r"""SELECT tval FROM db_dbsetting WHERE key = 'schema_generation';""") + except ProgrammingError: + result = backend.execute_raw(r"""SELECT val FROM db_dbsetting WHERE key = 'schema_generation';""") + + try: + return str(int(result[0][0])) + except (IndexError, TypeError, ValueError): + return '1' + + def get_schema_version_database(self): + """Return the database schema version. + + :return: `distutils.version.StrictVersion` with schema version of the database + """ + from django.db.utils import ProgrammingError + from aiida.manage.manager import get_manager + + backend = get_manager()._load_backend(schema_check=False) # pylint: disable=protected-access + + try: + result = backend.execute_raw(r"""SELECT tval FROM db_dbsetting WHERE key = 'db|schemaversion';""") + except ProgrammingError: + result = backend.execute_raw(r"""SELECT val FROM db_dbsetting WHERE key = 'db|schemaversion';""") + return result[0][0] + + def set_schema_version_database(self, version): + """Set the database schema version. + + :param version: string with schema version to set + """ + return self.get_settings_manager().set(SCHEMA_VERSION_KEY, version, description=SCHEMA_VERSION_DESCRIPTION) + + def _migrate_database_generation(self): + """Reset the database schema generation. + + For Django we also have to clear the `django_migrations` table that contains a history of all applied + migrations. After clearing it, we reinsert the name of the new initial schema . + """ + # pylint: disable=cyclic-import + from aiida.manage.manager import get_manager + super(DjangoBackendManager, self)._migrate_database_generation() + + backend = get_manager()._load_backend(schema_check=False) # pylint: disable=protected-access + backend.execute_raw(r"""DELETE FROM django_migrations WHERE app = 'db';""") + backend.execute_raw( + r"""INSERT INTO django_migrations (app, name, applied) VALUES ('db', '0001_initial', NOW());""" + ) + + def _migrate_database_version(self): + """Migrate the database to the current schema version.""" + super(DjangoBackendManager, self)._migrate_database_version() + from django.core.management import call_command # pylint: disable=no-name-in-module,import-error + call_command('migrate') + + +class DjangoSettingsManager(SettingsManager): + """Class to get, set and delete settings from the `DbSettings` table.""" + + table_name = 'db_dbsetting' + + 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 django.db import connection + if self.table_name not in connection.introspection.table_names(): + raise NotExistent('the settings table does not exist') + + def get(self, key): + """Return the setting with the given key. + + :param key: the key identifying the setting + :return: Setting + :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist + """ + from aiida.backends.djsite.db.models import DbSetting + + self.validate_table_existence() + setting = DbSetting.objects.filter(key=key).first() + + if setting is None: + raise NotExistent('setting `{}` does not exist'.format(key)) + + return Setting(setting.key, setting.val, setting.description, setting.time) + + def set(self, key, value, description=None): + """Return the settings with the given key. + + :param key: the key identifying the setting + :param value: the value for the setting + :param description: optional setting description + """ + from aiida.backends.djsite.db.models import DbSetting + from aiida.orm.utils.node import validate_attribute_extra_key + + self.validate_table_existence() + validate_attribute_extra_key(key) + + other_attribs = dict() + if description is not None: + other_attribs['description'] = description + + DbSetting.set_value(key, value, other_attribs=other_attribs) + + def delete(self, key): + """Delete the setting with the given key. + + :param key: the key identifying the setting + :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist + """ + from aiida.backends.djsite.db.models import DbSetting + + self.validate_table_existence() + + try: + DbSetting.del_value(key=key) + except KeyError: + raise NotExistent('setting `{}` does not exist'.format(key)) diff --git a/aiida/backends/djsite/utils.py b/aiida/backends/djsite/utils.py index fbd285e591..4f0b5b50a7 100644 --- a/aiida/backends/djsite/utils.py +++ b/aiida/backends/djsite/utils.py @@ -7,217 +7,23 @@ # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### -# pylint: disable=no-name-in-module,no-member,import-error,cyclic-import """Utility functions specific to the Django backend.""" from __future__ import division from __future__ import print_function from __future__ import absolute_import -import os -import django - -from aiida.backends.utils import validate_attribute_key, SettingsManager, Setting, validate_schema_generation -from aiida.common import NotExistent - -SCHEMA_VERSION_DB_KEY = 'db|schemaversion' -SCHEMA_VERSION_DB_DESCRIPTION = 'The version of the schema used in this database.' - - -class DjangoSettingsManager(SettingsManager): - """Class to get, set and delete settings from the `DbSettings` table.""" - - table_name = 'db_dbsetting' - - 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 django.db import connection - if self.table_name not in connection.introspection.table_names(): - raise NotExistent('the settings table does not exist') - - def get(self, key): - """Return the setting with the given key. - - :param key: the key identifying the setting - :return: Setting - :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist - """ - from aiida.backends.djsite.db.models import DbSetting - - self.validate_table_existence() - setting = DbSetting.objects.filter(key=key).first() - - if setting is None: - raise NotExistent('setting `{}` does not exist'.format(key)) - - return Setting(setting.key, setting.val, setting.description, setting.time) - - def set(self, key, value, description=None): - """Return the settings with the given key. - - :param key: the key identifying the setting - :param value: the value for the setting - :param description: optional setting description - """ - from aiida.backends.djsite.db.models import DbSetting - - self.validate_table_existence() - validate_attribute_key(key) - - other_attribs = dict() - if description is not None: - other_attribs['description'] = description - - DbSetting.set_value(key, value, other_attribs=other_attribs) - - def delete(self, key): - """Delete the setting with the given key. - - :param key: the key identifying the setting - :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist - """ - from aiida.backends.djsite.db.models import DbSetting - - self.validate_table_existence() - - try: - DbSetting.del_value(key=key) - except KeyError: - raise NotExistent('setting `{}` does not exist'.format(key)) - - -def load_dbenv(profile): - """Load the database environment and ensure that the code and database schema versions are compatible. - - :param profile: instance of `Profile` whose database to load - """ - _load_dbenv_noschemacheck(profile) - check_schema_version(profile_name=profile.name) - - -def _load_dbenv_noschemacheck(profile): # pylint: disable=unused-argument - """Load the database environment without checking that code and database schema versions are compatible. - - This should ONLY be used internally, inside load_dbenv, and for schema migrations. DO NOT USE OTHERWISE! - - :param profile: instance of `Profile` whose database to load - """ - os.environ['DJANGO_SETTINGS_MODULE'] = 'aiida.backends.djsite.settings' - django.setup() - - -def unload_dbenv(): - """Unload the database environment. - - This means that the settings in `aiida.backends.djsite.settings` are "unset". - This needs to implemented and will address #2813 - """ - - -_aiida_autouser_cache = None # pylint: disable=invalid-name - - -def migrate_database(): - """Migrate the database to the latest schema version.""" - validate_schema_generation() - from django.core.management import call_command - call_command('migrate') - - -def check_schema_version(profile_name): - """ - Check if the version stored in the database is the same of the version - of the code. - - :note: if the DbSetting table does not exist, this function does not - fail. The reason is to avoid to have problems before running the first - migrate call. - - :note: if no version is found, the version is set to the version of the - code. This is useful to have the code automatically set the DB version - at the first code execution. - - :raise aiida.common.ConfigurationError: if the two schema versions do not match. - Otherwise, just return. - """ - # pylint: disable=duplicate-string-formatting-argument - from django.db import connection - - import aiida.backends.djsite.db.models - from aiida.common.exceptions import ConfigurationError - - # Do not do anything if the table does not exist yet - if 'db_dbsetting' not in connection.introspection.table_names(): - return - - validate_schema_generation() - - code_schema_version = aiida.backends.djsite.db.models.SCHEMA_VERSION - db_schema_version = get_db_schema_version() - - if db_schema_version is None: - # No code schema defined yet, I set it to the code version - set_db_schema_version(code_schema_version) - db_schema_version = get_db_schema_version() - - if code_schema_version > db_schema_version: - raise ConfigurationError( - 'Database schema version {} is outdated compared to the code schema version {}.\n' - 'Note: Have all calculations and workflows have finished running? ' - 'If not, revert the code to the previous version and let them finish before upgrading.\n' - 'To migrate the database to the current code version, run the following commands:' - '\n verdi -p {} daemon stop\n verdi -p {} database migrate'.format( - db_schema_version, code_schema_version, profile_name, profile_name - ) - ) - elif code_schema_version < db_schema_version: - raise ConfigurationError( - 'Database schema version {} is newer than code schema version {}.\n' - 'You cannot use an outdated code with a newer database. Please upgrade.'.format( - db_schema_version, code_schema_version - ) - ) - - -def set_db_schema_version(version): - """ - Set the schema version stored in the DB. Use only if you know what you are doing. - """ - from aiida.backends.utils import get_settings_manager - manager = get_settings_manager() - return manager.set(SCHEMA_VERSION_DB_KEY, version, description=SCHEMA_VERSION_DB_DESCRIPTION) - - -def get_db_schema_version(): - """ - Get the current schema version stored in the DB. Return None if - it is not stored. - """ - from django.db.utils import ProgrammingError - from aiida.manage.manager import get_manager - backend = get_manager()._load_backend(schema_check=False) # pylint: disable=protected-access - - try: - result = backend.execute_raw(r"""SELECT tval FROM db_dbsetting WHERE key = 'db|schemaversion';""") - except ProgrammingError: - result = backend.execute_raw(r"""SELECT val FROM db_dbsetting WHERE key = 'db|schemaversion';""") - - return result[0][0] - def delete_nodes_and_connections_django(pks_to_delete): # pylint: disable=invalid-name - """ - Delete all nodes corresponding to pks in the input. + """Delete all nodes corresponding to pks in the input. + :param pks_to_delete: A list, tuple or set of pks that should be deleted. """ + # pylint: disable=no-member,import-error,no-name-in-module from django.db import transaction from django.db.models import Q from aiida.backends.djsite.db import models with transaction.atomic(): # This is fixed in pylint-django>=2, but this supports only py3 - # pylint: disable=no-member # Delete all links pointing to or from a given node models.DbLink.objects.filter(Q(input__in=pks_to_delete) | Q(output__in=pks_to_delete)).delete() # now delete nodes diff --git a/aiida/backends/manager.py b/aiida/backends/manager.py new file mode 100644 index 0000000000..7497fbe131 --- /dev/null +++ b/aiida/backends/manager.py @@ -0,0 +1,313 @@ +# -*- 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 # +########################################################################### +"""Module for settings and utilities to determine and set the database schema versions.""" +from __future__ import absolute_import + +import abc +import collections + +from aiida.common import exceptions + +SCHEMA_VERSION_KEY = 'db|schemaversion' +SCHEMA_VERSION_DESCRIPTION = 'Database schema version' + +SCHEMA_GENERATION_KEY = 'schema_generation' # The key to store the database schema generation in the settings table +SCHEMA_GENERATION_DESCRIPTION = 'Database schema generation' +SCHEMA_GENERATION_VALUE = '1' # The current schema generation + +# Mapping of schema generation onto a tuple of valid schema reset generation and `aiida-core` version number. Given the +# current code schema generation as the key, the first element of the tuple tells what schema generation the database +# should have to be able to reset the schema. If the generation of the database is correct, but the schema version of +# the database does not match the one required for the reset, it means the user first has to downgrade the `aiida-core` +# version and perform the latest migrations. The required version is provided by the tuples second element. +SCHEMA_GENERATION_RESET = { + '1': ('1', '1.*'), +} + +TEMPLATE_INVALID_SCHEMA_GENERATION = """ +Database schema generation `{schema_generation_database}` is incompatible with the required schema generation `{schema_generation_code}`. +To migrate the database schema generation to the current one, run the following command: + + verdi -p {profile_name} database migrate +""" + +TEMPLATE_INVALID_SCHEMA_VERSION = """ +Database schema version `{schema_version_database}` is incompatible with the required schema version `{schema_version_code}`. +To migrate the database schema version to the current one, run the following command: + + verdi -p {profile_name} database migrate +""" + +TEMPLATE_MIGRATE_SCHEMA_VERSION_INVALID_VERSION = """ +Cannot migrate the database version from `{schema_version_database}` to `{schema_version_code}`. +The database version is ahead of the version of the code and downgrades of the database are not supported. +""" + +TEMPLATE_MIGRATE_SCHEMA_GENERATION_INVALID_GENERATION = """ +Cannot migrate database schema generation from `{schema_generation_database}` to `{schema_generation_code}`. +This version of `aiida-core` can only migrate databases with schema generation `{schema_generation_reset}` +""" + +TEMPLATE_MIGRATE_SCHEMA_GENERATION_INVALID_VERSION = """ +Cannot migrate database schema generation from `{schema_generation_database}` to `{schema_generation_code}`. +The current database version is `{schema_version_database}` but `{schema_version_reset}` is required for generation migration. +First install `aiida-core~={aiida_core_version_reset}` and migrate the database to the latest version. +After the database schema is migrated to version `{schema_version_reset}` you can reinstall this version of `aiida-core` and migrate the schema generation. +""" + +Setting = collections.namedtuple('Setting', ['key', 'value', 'description', 'time']) + + +class SettingsManager(object): # pylint: disable=useless-object-inheritance + """Class to get, set and delete settings from the `DbSettings` table.""" + + @abc.abstractmethod + def get(self, key): + """Return the setting with the given key. + + :param key: the key identifying the setting + :return: Setting + :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist + """ + + @abc.abstractmethod + def set(self, key, value, description=None): + """Return the settings with the given key. + + :param key: the key identifying the setting + :param value: the value for the setting + :param description: optional setting description + """ + + @abc.abstractmethod + def delete(self, key): + """Delete the setting with the given key. + + :param key: the key identifying the setting + :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist + """ + + +class BackendManager(object): # pylint: disable=useless-object-inheritance + """Class to manage the database schema and environment.""" + + _settings_manager = None + + @abc.abstractmethod + def get_settings_manager(self): + """Return an instance of the `SettingsManager`. + + :return: `SettingsManager` + """ + + def load_backend_environment(self, profile, validate_schema=True): + """Load the backend environment. + + :param profile: the profile whose backend environment to load + :param validate_schema: boolean, if True, validate the schema first before loading the environment. + """ + self._load_backend_environment() + + if validate_schema: + self.validate_schema(profile) + + @abc.abstractmethod + def _load_backend_environment(self): + """Load the backend environment.""" + + @abc.abstractmethod + def reset_backend_environment(self): + """Reset the backend environment.""" + + def migrate(self): + """Migrate the database to the latest schema generation or version.""" + try: + # If the settings table does not exist, we are dealing with an empty database. We cannot perform the checks + # because they rely on the settings table existing, so instead we do not validate but directly call method + # `_migrate_database_version` which will perform the migration to create the initial schema. + self.get_settings_manager().validate_table_existence() + except exceptions.NotExistent: + self._migrate_database_version() + return + + if SCHEMA_GENERATION_VALUE != self.get_schema_generation_database(): + self.validate_schema_generation_for_migration() + self._migrate_database_generation() + + if self.get_schema_version_code() != self.get_schema_version_database(): + self.validate_schema_version_for_migration() + self._migrate_database_version() + + def _migrate_database_generation(self): + """Migrate the database schema generation. + + .. warning:: this should NEVER be called directly because there is no validation performed on whether the + current database schema generation and version can actually be migrated. + + This normally just consists out of setting the schema generation value, but depending on the backend more might + be needed. In that case, this method should be overridden and call `super` first, followed by the additional + logic that is required. + """ + self.set_schema_generation_database(SCHEMA_GENERATION_VALUE) + self.set_schema_version_database(self.get_schema_version_code()) + + def _migrate_database_version(self): + """Migrate the database to the current schema version. + + .. warning:: this should NEVER be called directly because there is no validation performed on whether the + current database schema generation and version can actually be migrated. + """ + + @abc.abstractmethod + def is_database_schema_ahead(self): + """Determine whether the database schema version is ahead of the code schema version. + + .. warning:: this will not check whether the schema generations are equal + + :return: boolean, True if the database schema version is ahead of the code schema version. + """ + + @abc.abstractmethod + def get_schema_version_code(self): + """Return the code schema version.""" + + @abc.abstractmethod + def get_schema_version_reset(self, schema_generation_code): + """Return schema version the database should have to be able to automatically reset to code schema generation. + + :param schema_generation_code: the schema generation of the code. + :return: schema version + """ + + @abc.abstractmethod + def get_schema_version_database(self): + """Return the database schema version. + + :return: `distutils.version.LooseVersion` with schema version of the database + """ + + @abc.abstractmethod + def set_schema_version_database(self, version): + """Set the database schema version. + + :param version: string with schema version to set + """ + + def get_schema_generation_database(self): + """Return the database schema generation. + + :return: `distutils.version.LooseVersion` with schema generation of the database + """ + try: + setting = self.get_settings_manager().get(SCHEMA_GENERATION_KEY) + return setting.value + except exceptions.NotExistent: + return '1' + + def set_schema_generation_database(self, generation): + """Set the database schema generation. + + :param generation: string with schema generation to set + """ + self.get_settings_manager().set(SCHEMA_GENERATION_KEY, generation) + + def validate_schema(self, profile): + """Validate that the current database generation and schema are up-to-date with that of the code. + + :param profile: the profile for which to validate the database schema + :raises `aiida.common.exceptions.ConfigurationError`: if database schema version or generation is not up-to-date + """ + self.validate_schema_generation(profile) + self.validate_schema_version(profile) + + def validate_schema_generation_for_migration(self): + """Validate whether the current database schema generation can be migrated. + + :raises `aiida.common.exceptions.IncompatibleDatabaseSchema`: if database schema generation cannot be migrated + """ + schema_generation_code = SCHEMA_GENERATION_VALUE + schema_generation_database = self.get_schema_generation_database() + schema_version_database = self.get_schema_version_database() + schema_version_reset = self.get_schema_version_reset(schema_generation_code) + schema_generation_reset, aiida_core_version_reset = SCHEMA_GENERATION_RESET[schema_generation_code] + + if schema_generation_database != schema_generation_reset: + raise exceptions.IncompatibleDatabaseSchema( + TEMPLATE_MIGRATE_SCHEMA_GENERATION_INVALID_GENERATION.format( + schema_generation_database=schema_generation_database, + schema_generation_code=schema_generation_code, + schema_generation_reset=schema_generation_reset + ) + ) + + if schema_version_database != schema_version_reset: + raise exceptions.IncompatibleDatabaseSchema( + TEMPLATE_MIGRATE_SCHEMA_GENERATION_INVALID_VERSION.format( + schema_generation_database=schema_generation_database, + schema_generation_code=schema_generation_code, + schema_version_database=schema_version_database, + schema_version_reset=schema_version_reset, + aiida_core_version_reset=aiida_core_version_reset + ) + ) + + def validate_schema_version_for_migration(self): + """Validate whether the current database schema version can be migrated. + + .. warning:: this will not validate that the schema generation is correct. + + :raises `aiida.common.exceptions.IncompatibleDatabaseSchema`: if database schema version cannot be migrated + """ + schema_version_code = self.get_schema_version_code() + schema_version_database = self.get_schema_version_database() + + if self.is_database_schema_ahead(): + # Database is newer than the code so a downgrade would be necessary but this is not supported. + raise exceptions.IncompatibleDatabaseSchema( + TEMPLATE_MIGRATE_SCHEMA_VERSION_INVALID_VERSION.format( + schema_version_database=schema_version_database, + schema_version_code=schema_version_code, + ) + ) + + def validate_schema_generation(self, profile): + """Validate that the current database schema generation is up-to-date with that of the code. + + :raises `aiida.common.exceptions.IncompatibleDatabaseSchema`: if database schema generation is not up-to-date + """ + schema_generation_code = SCHEMA_GENERATION_VALUE + schema_generation_database = self.get_schema_generation_database() + + if schema_generation_database != schema_generation_code: + raise exceptions.IncompatibleDatabaseSchema( + TEMPLATE_INVALID_SCHEMA_GENERATION.format( + schema_generation_database=schema_generation_database, + schema_generation_code=schema_generation_code, + profile_name=profile.name, + ) + ) + + def validate_schema_version(self, profile): + """Validate that the current database schema version is up-to-date with that of the code. + + :param profile: the profile for which to validate the database schema + :raises `aiida.common.exceptions.IncompatibleDatabaseSchema`: if database schema version is not up-to-date + """ + schema_version_code = self.get_schema_version_code() + schema_version_database = self.get_schema_version_database() + + if schema_version_database != schema_version_code: + raise exceptions.IncompatibleDatabaseSchema( + TEMPLATE_INVALID_SCHEMA_VERSION.format( + schema_version_database=schema_version_database, + schema_version_code=schema_version_code, + profile_name=profile.name + ) + ) diff --git a/aiida/backends/sqlalchemy/__init__.py b/aiida/backends/sqlalchemy/__init__.py index 8b3365a419..d5b04d89f2 100644 --- a/aiida/backends/sqlalchemy/__init__.py +++ b/aiida/backends/sqlalchemy/__init__.py @@ -7,6 +7,8 @@ # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### +# pylint: disable=import-error,no-name-in-module,global-statement +"""Module with implementation of the database backend using SqlAlchemy.""" from __future__ import division from __future__ import print_function from __future__ import absolute_import @@ -34,7 +36,7 @@ def get_scoped_session(): return SCOPED_SESSION_CLASS() -def recreate_after_fork(engine): +def recreate_after_fork(engine): # pylint: disable=unused-argument """Callback called after a fork. Not only disposes the engine, but also recreates a new scoped session to use independent sessions in the fork. @@ -72,7 +74,8 @@ def reset_session(profile=None): password=profile.database_password, hostname=profile.database_hostname, port=profile.database_port, - name=profile.database_name) + name=profile.database_name + ) ENGINE = create_engine(engine_url, json_serializer=json.dumps, json_deserializer=json.loads, encoding='utf-8') SCOPED_SESSION_CLASS = scoped_session(sessionmaker(bind=ENGINE, expire_on_commit=True)) diff --git a/aiida/backends/sqlalchemy/manage.py b/aiida/backends/sqlalchemy/manage.py index 049e7f769d..9e4780da23 100755 --- a/aiida/backends/sqlalchemy/manage.py +++ b/aiida/backends/sqlalchemy/manage.py @@ -8,7 +8,6 @@ # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### -# pylint: """Simple wrapper around the alembic command line tool that first loads an AiiDA profile.""" from __future__ import division from __future__ import print_function @@ -26,14 +25,13 @@ def execute_alembic_command(command_name, **kwargs): :param command_name: the sub command name :param kwargs: parameters to pass to the command """ - from aiida.backends import sqlalchemy as sa - from aiida.backends.sqlalchemy.utils import get_alembic_conf + from aiida.backends.sqlalchemy.manager import SqlaBackendManager - with sa.ENGINE.begin() as connection: - alembic_cfg = get_alembic_conf() - alembic_cfg.attributes['connection'] = connection # pylint: disable=unsupported-assignment-operation + manager = SqlaBackendManager() + + with manager.alembic_config() as config: command = getattr(alembic.command, command_name) - command(alembic_cfg, **kwargs) + command(config, **kwargs) @click.group() diff --git a/aiida/backends/sqlalchemy/manager.py b/aiida/backends/sqlalchemy/manager.py new file mode 100644 index 0000000000..625b0b6e88 --- /dev/null +++ b/aiida/backends/sqlalchemy/manager.py @@ -0,0 +1,210 @@ +# -*- coding: utf-8 -*- +# pylint: disable=import-error,no-name-in-module +"""Utilities and configuration of the SqlAlchemy database schema.""" +from __future__ import absolute_import + +import os +import contextlib + +from alembic import command +from alembic.config import Config +from alembic.runtime.environment import EnvironmentContext +from alembic.script import ScriptDirectory + +from sqlalchemy.orm.exc import NoResultFound + +from aiida.backends.sqlalchemy import get_scoped_session +from aiida.common import NotExistent +from ..manager import BackendManager, SettingsManager, Setting + +ALEMBIC_FILENAME = 'alembic.ini' +ALEMBIC_REL_PATH = 'migrations' + +# The database schema version required to perform schema reset for a given code schema generation +SCHEMA_VERSION_RESET = {'1': None} + + +class SqlaBackendManager(BackendManager): + """Class to manage the database schema.""" + + @staticmethod + @contextlib.contextmanager + def alembic_config(): + """Context manager to return an instance of an Alembic configuration with the current connection inserted. + + :return: instance of :py:class:`alembic.config.Config` + """ + from . import ENGINE + + with ENGINE.begin() as connection: + dir_path = os.path.dirname(os.path.realpath(__file__)) + config = Config(os.path.join(dir_path, ALEMBIC_FILENAME)) + config.set_main_option('script_location', os.path.join(dir_path, ALEMBIC_REL_PATH)) + config.attributes['connection'] = connection # pylint: disable=unsupported-assignment-operation + yield config + + def get_settings_manager(self): + """Return an instance of the `SettingsManager`. + + :return: `SettingsManager` + """ + if self._settings_manager is None: + self._settings_manager = SqlaSettingsManager() + + return self._settings_manager + + def _load_backend_environment(self): + """Load the backend environment.""" + from . import reset_session + reset_session() + + def reset_backend_environment(self): + """Reset the backend environment.""" + from aiida.backends import sqlalchemy + if sqlalchemy.ENGINE is not None: + sqlalchemy.ENGINE.dispose() + sqlalchemy.SCOPED_SESSION_CLASS = None + + def is_database_schema_ahead(self): + """Determine whether the database schema version is ahead of the code schema version. + + .. warning:: this will not check whether the schema generations are equal + + :return: boolean, True if the database schema version is ahead of the code schema version. + """ + from alembic.util import CommandError + + # In the case of SqlAlchemy, if the database revision is ahead of the code, that means the revision stored in + # the database is not even present in the code base. Therefore we cannot locate it in the revision graph and + # determine whether it is ahead of the current code head. We simply try to get the revision and if it does not + # exist it means it is ahead. + with self.alembic_config() as config: + try: + script = ScriptDirectory.from_config(config) + script.get_revision(self.get_schema_version_database()) + except CommandError: + # Raised when the revision of the database is not present in the revision graph. + return True + else: + return False + + def get_schema_version_code(self): + """Return the code schema version.""" + with self.alembic_config() as config: + script = ScriptDirectory.from_config(config) + schema_version_code = script.get_current_head() + + return schema_version_code + + def get_schema_version_reset(self, schema_generation_code): + """Return schema version the database should have to be able to automatically reset to code schema generation. + + :param schema_generation_code: the schema generation of the code. + :return: schema version + """ + return SCHEMA_VERSION_RESET[schema_generation_code] + + def get_schema_version_database(self): + """Return the database schema version. + + :return: `distutils.version.StrictVersion` with schema version of the database + """ + + def get_database_version(revision, _): + """Get the current revision.""" + if isinstance(revision, tuple) and revision: + config.attributes['rev'] = revision[0] # pylint: disable=unsupported-assignment-operation + else: + config.attributes['rev'] = None # pylint: disable=unsupported-assignment-operation + return [] + + with self.alembic_config() as config: + + script = ScriptDirectory.from_config(config) + + with EnvironmentContext(config, script, fn=get_database_version): + script.run_env() + return config.attributes['rev'] # pylint: disable=unsubscriptable-object + + def set_schema_version_database(self, version): + """Set the database schema version. + + :param version: string with schema version to set + """ + # pylint: disable=cyclic-import + from aiida.manage.manager import get_manager + backend = get_manager()._load_backend(schema_check=False) # pylint: disable=protected-access + backend.execute_raw(r"""UPDATE alembic_version SET version_num='{}';""".format(version)) + + def _migrate_database_version(self): + """Migrate the database to the current schema version.""" + super(SqlaBackendManager, self)._migrate_database_version() + with self.alembic_config() as config: + command.upgrade(config, 'head') + + +class SqlaSettingsManager(SettingsManager): + """Class to get, set and delete settings from the `DbSettings` table.""" + + table_name = 'db_dbsetting' + + 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 get(self, key): + """Return the setting with the given key. + + :param key: the key identifying the setting + :return: Setting + :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist + """ + from aiida.backends.sqlalchemy.models.settings import DbSetting + self.validate_table_existence() + + try: + setting = get_scoped_session().query(DbSetting).filter_by(key=key).one() + except NoResultFound: + raise NotExistent('setting `{}` does not exist'.format(key)) + + return Setting(key, setting.getvalue(), setting.description, setting.time) + + def set(self, key, value, description=None): + """Return the settings with the given key. + + :param key: the key identifying the setting + :param value: the value for the setting + :param description: optional setting description + """ + from aiida.backends.sqlalchemy.models.settings import DbSetting + from aiida.orm.utils.node import validate_attribute_extra_key + + self.validate_table_existence() + validate_attribute_extra_key(key) + + other_attribs = dict() + if description is not None: + other_attribs['description'] = description + + DbSetting.set_value(key, value, other_attribs=other_attribs) + + def delete(self, key): + """Delete the setting with the given key. + + :param key: the key identifying the setting + :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist + """ + from aiida.backends.sqlalchemy.models.settings import DbSetting + self.validate_table_existence() + + try: + setting = get_scoped_session().query(DbSetting).filter_by(key=key).one() + setting.delete() + except NoResultFound: + raise NotExistent('setting `{}` does not exist'.format(key)) diff --git a/aiida/backends/sqlalchemy/models/settings.py b/aiida/backends/sqlalchemy/models/settings.py index 14eeb99a55..4f391a76fa 100644 --- a/aiida/backends/sqlalchemy/models/settings.py +++ b/aiida/backends/sqlalchemy/models/settings.py @@ -15,12 +15,12 @@ from sqlalchemy import Column from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.orm.attributes import flag_modified from sqlalchemy.schema import UniqueConstraint from sqlalchemy.types import Integer, String, DateTime from aiida.backends import sqlalchemy as sa from aiida.backends.sqlalchemy.models.base import Base -from aiida.backends.sqlalchemy.utils import flag_modified from aiida.common import timezone diff --git a/aiida/backends/sqlalchemy/queries.py b/aiida/backends/sqlalchemy/queries.py index 033a3c6b88..649b31e465 100644 --- a/aiida/backends/sqlalchemy/queries.py +++ b/aiida/backends/sqlalchemy/queries.py @@ -11,8 +11,6 @@ from __future__ import print_function from __future__ import absolute_import -from contextlib import contextmanager - from aiida.backends.general.abstractqueries import AbstractQueryManager diff --git a/aiida/backends/sqlalchemy/tests/test_migrations.py b/aiida/backends/sqlalchemy/tests/test_migrations.py index a26cc78029..b083fd88c7 100644 --- a/aiida/backends/sqlalchemy/tests/test_migrations.py +++ b/aiida/backends/sqlalchemy/tests/test_migrations.py @@ -23,7 +23,7 @@ from aiida.backends import sqlalchemy as sa from aiida.backends.general.migrations import utils -from aiida.backends.sqlalchemy import utils as sqlalchemy_utils +from aiida.backends.sqlalchemy import manager from aiida.backends.sqlalchemy.models.base import Base from aiida.backends.sqlalchemy.tests.test_utils import new_database from aiida.backends.sqlalchemy.utils import flag_modified @@ -52,10 +52,7 @@ def setUpClass(cls, *args, **kwargs): Prepare the test class with the alembivc configuration """ super(TestMigrationsSQLA, cls).setUpClass(*args, **kwargs) - cls.migr_method_dir_path = os.path.dirname(os.path.realpath(sqlalchemy_utils.__file__)) - alembic_dpath = os.path.join(cls.migr_method_dir_path, sqlalchemy_utils.ALEMBIC_REL_PATH) - cls.alembic_cfg = Config() - cls.alembic_cfg.set_main_option('script_location', alembic_dpath) + cls.manager = manager.SqlaBackendManager() def setUp(self): """ @@ -88,9 +85,8 @@ def migrate_db_up(self, destination): :param destination: the name of the destination migration """ # Undo all previous real migration of the database - with sa.ENGINE.connect() as connection: - self.alembic_cfg.attributes['connection'] = connection # pylint: disable=unsupported-assignment-operation - command.upgrade(self.alembic_cfg, destination) + with self.manager.alembic_config() as config: + command.upgrade(config, destination) def migrate_db_down(self, destination): """ @@ -98,9 +94,8 @@ def migrate_db_down(self, destination): :param destination: the name of the destination migration """ - with sa.ENGINE.connect() as connection: - self.alembic_cfg.attributes['connection'] = connection # pylint: disable=unsupported-assignment-operation - command.downgrade(self.alembic_cfg, destination) + with self.manager.alembic_config() as config: + command.downgrade(config, destination) def tearDown(self): """ @@ -291,9 +286,9 @@ def setUp(self): from sqlalchemydiff.util import get_temporary_uri from aiida.backends.sqlalchemy.migrations import versions - self.migr_method_dir_path = os.path.dirname(os.path.realpath(sqlalchemy_utils.__file__)) + self.migr_method_dir_path = os.path.dirname(os.path.realpath(manager.__file__)) # Set the alembic script directory location - self.alembic_dpath = os.path.join(self.migr_method_dir_path, sqlalchemy_utils.ALEMBIC_REL_PATH) + self.alembic_dpath = os.path.join(self.migr_method_dir_path, manager.ALEMBIC_REL_PATH) # pylint: disable=no-member # Constructing the versions directory versions_dpath = os.path.join(os.path.dirname(versions.__file__)) diff --git a/aiida/backends/sqlalchemy/utils.py b/aiida/backends/sqlalchemy/utils.py index 6ac58008ef..a2e2b64e5c 100644 --- a/aiida/backends/sqlalchemy/utils.py +++ b/aiida/backends/sqlalchemy/utils.py @@ -7,91 +7,39 @@ # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### +# pylint: disable=import-error,no-name-in-module +"""Utility functions specific to the SqlAlchemy backend.""" from __future__ import division from __future__ import absolute_import from __future__ import print_function -from alembic import command -from alembic.config import Config -from alembic.runtime.environment import EnvironmentContext -from alembic.script import ScriptDirectory -from sqlalchemy.orm.exc import NoResultFound - -from aiida.backends import sqlalchemy as sa -from aiida.backends.sqlalchemy import get_scoped_session -from aiida.backends.utils import validate_attribute_key, SettingsManager, Setting, validate_schema_generation -from aiida.common import NotExistent - - -ALEMBIC_FILENAME = 'alembic.ini' -ALEMBIC_REL_PATH = 'migrations' - - -class SqlaSettingsManager(SettingsManager): - """Class to get, set and delete settings from the `DbSettings` table.""" - - table_name = 'db_dbsetting' - - 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 get(self, key): - """Return the setting with the given key. - - :param key: the key identifying the setting - :return: Setting - :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist - """ - from aiida.backends.sqlalchemy.models.settings import DbSetting - self.validate_table_existence() - - try: - setting = get_scoped_session().query(DbSetting).filter_by(key=key).one() - except NoResultFound: - raise NotExistent('setting `{}` does not exist'.format(key)) - - return Setting(key, setting.getvalue(), setting.description, setting.time) - - def set(self, key, value, description=None): - """Return the settings with the given key. - - :param key: the key identifying the setting - :param value: the value for the setting - :param description: optional setting description - """ - from aiida.backends.sqlalchemy.models.settings import DbSetting - self.validate_table_existence() - validate_attribute_key(key) - - other_attribs = dict() - if description is not None: - other_attribs['description'] = description - - DbSetting.set_value(key, value, other_attribs=other_attribs) - - def delete(self, key): - """Delete the setting with the given key. - - :param key: the key identifying the setting - :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist - """ - from aiida.backends.sqlalchemy.models.settings import DbSetting - self.validate_table_existence() +def delete_nodes_and_connections_sqla(pks_to_delete): # pylint: disable=invalid-name + """ + Delete all nodes corresponding to pks in the input. + :param pks_to_delete: A list, tuple or set of pks that should be deleted. + """ + # pylint: disable=no-value-for-parameter + from aiida.backends.sqlalchemy.models.node import DbNode, DbLink + from aiida.backends.sqlalchemy.models.group import table_groups_nodes + from aiida.manage.manager import get_manager - try: - setting = get_scoped_session().query(DbSetting).filter_by(key=key).one() - setting.delete() - except NoResultFound: - raise NotExistent('setting `{}` does not exist'.format(key)) + backend = get_manager().get_backend() + with backend.transaction() as session: + # I am first making a statement to delete the membership of these nodes to groups. + # Since table_groups_nodes is a sqlalchemy.schema.Table, I am using expression language to compile + # a stmt to be executed by the session. It works, but it's not nice that two different ways are used! + # Can this be changed? + stmt = table_groups_nodes.delete().where(table_groups_nodes.c.dbnode_id.in_(list(pks_to_delete))) + session.execute(stmt) + # First delete links, then the Nodes, since we are not cascading deletions. + # Here I delete the links coming out of the nodes marked for deletion. + session.query(DbLink).filter(DbLink.input_id.in_(list(pks_to_delete))).delete(synchronize_session='fetch') + # Here I delete the links pointing to the nodes marked for deletion. + session.query(DbLink).filter(DbLink.output_id.in_(list(pks_to_delete))).delete(synchronize_session='fetch') + # Now I am deleting the nodes + session.query(DbNode).filter(DbNode.id.in_(list(pks_to_delete))).delete(synchronize_session='fetch') def flag_modified(instance, key): @@ -106,42 +54,11 @@ def flag_modified(instance, key): from aiida.orm.implementation.sqlalchemy.utils import ModelWrapper if isinstance(instance, ModelWrapper): - instance = instance._model + instance = instance._model # pylint: disable=protected-access flag_modified_sqla(instance, key) -def load_dbenv(profile): - """Load the database environment and ensure that the code and database schema versions are compatible. - - :param profile: the string with the profile to use - """ - _load_dbenv_noschemacheck(profile) - check_schema_version(profile_name=profile.name) - - -def _load_dbenv_noschemacheck(profile): - """Load the database environment without checking that code and database schema versions are compatible. - - This should ONLY be used internally, inside load_dbenv, and for schema migrations. DO NOT USE OTHERWISE! - - :param profile: instance of `Profile` whose database to load - """ - sa.reset_session(profile) - - -def unload_dbenv(): - """Unload the database environment, which boils down to destroying the current engine and session.""" - if sa.ENGINE is not None: - sa.ENGINE.dispose() - sa.SCOPED_SESSION_CLASS = None - - -_aiida_autouser_cache = None - - -# XXX the code here isn't different from the one use in Django. We may be able -# to refactor it in some way def install_tc(session): """ Install the transitive closure table with SqlAlchemy. @@ -153,24 +70,25 @@ def install_tc(session): closure_table_parent_field = 'parent_id' closure_table_child_field = 'child_id' - session.execute(get_pg_tc(links_table_name, links_table_input_field, - links_table_output_field, closure_table_name, - closure_table_parent_field, - closure_table_child_field)) + session.execute( + get_pg_tc( + links_table_name, links_table_input_field, links_table_output_field, closure_table_name, + closure_table_parent_field, closure_table_child_field + ) + ) -def get_pg_tc(links_table_name, - links_table_input_field, - links_table_output_field, - closure_table_name, - closure_table_parent_field, - closure_table_child_field): +def get_pg_tc( + links_table_name, links_table_input_field, links_table_output_field, closure_table_name, closure_table_parent_field, + closure_table_child_field +): """ Return the transitive closure table template """ from string import Template - pg_tc = Template(""" + pg_tc = Template( + """ DROP TRIGGER IF EXISTS autoupdate_tc ON $links_table_name; DROP FUNCTION IF EXISTS update_tc(); @@ -333,140 +251,13 @@ def get_pg_tc(links_table_name, ON $links_table_name FOR each ROW EXECUTE PROCEDURE update_tc(); -""") - return pg_tc.substitute(links_table_name=links_table_name, - links_table_input_field=links_table_input_field, - links_table_output_field=links_table_output_field, - closure_table_name=closure_table_name, - closure_table_parent_field=closure_table_parent_field, - closure_table_child_field=closure_table_child_field) - - -def migrate_database(alembic_cfg=None): - """Migrate the database to the latest schema version. - - :param config: alembic configuration to use, will use default if not provided - """ - validate_schema_generation() - from aiida.backends import sqlalchemy as sa - - if alembic_cfg is None: - alembic_cfg = get_alembic_conf() - - with sa.ENGINE.connect() as connection: - alembic_cfg.attributes['connection'] = connection - command.upgrade(alembic_cfg, 'head') - - -def check_schema_version(profile_name): - """ - Check if the version stored in the database is the same of the version of the code. - - :raise aiida.common.ConfigurationError: if the two schema versions do not match - """ - from aiida.backends import sqlalchemy as sa - from aiida.common.exceptions import ConfigurationError - - alembic_cfg = get_alembic_conf() - - validate_schema_generation() - - # Getting the version of the code and the database, reusing the existing engine (initialized by AiiDA) - with sa.ENGINE.begin() as connection: - alembic_cfg.attributes['connection'] = connection - code_schema_version = get_migration_head(alembic_cfg) - db_schema_version = get_db_schema_version(alembic_cfg) - - if code_schema_version != db_schema_version: - raise ConfigurationError('Database schema version {} is outdated compared to the code schema version {}\n' - 'Before you upgrade, make sure all calculations and workflows have finished running.\n' - 'If this is not the case, revert the code to the previous version and finish them first.\n' - 'To migrate the database to the current version, run the following commands:' - '\n verdi -p {} daemon stop\n verdi -p {} database migrate'.format( - db_schema_version, code_schema_version, profile_name, profile_name)) - - - -def get_migration_head(config): - """ - This function returns the head of the migration scripts. - :param config: The alembic configuration. - :return: The version of the head. - """ - script = ScriptDirectory.from_config(config) - return script.get_current_head() - - -def get_db_schema_version(config): - """ - This function returns the current version of the database. - :param config: The alembic configuration. - :return: The version of the database. - """ - if config is None: - return None - - script = ScriptDirectory.from_config(config) - - def get_db_version(rev, _): - if isinstance(rev, tuple) and len(rev) > 0: - config.attributes['rev'] = rev[0] - else: - config.attributes['rev'] = None - - return [] - - with EnvironmentContext( - config, - script, - fn=get_db_version - ): - script.run_env() - return config.attributes['rev'] - - -def get_alembic_conf(): - """ - This function returns the alembic configuration file contents by doing - the necessary updates in the 'script_location' name. - :return: The alembic configuration. - """ - # Constructing the alembic full path & getting the configuration - import os - dir_path = os.path.dirname(os.path.realpath(__file__)) - alembic_fpath = os.path.join(dir_path, ALEMBIC_FILENAME) - alembic_cfg = Config(alembic_fpath) - - # Set the alembic script directory location - alembic_dpath = os.path.join(dir_path, ALEMBIC_REL_PATH) - alembic_cfg.set_main_option('script_location', alembic_dpath) - - return alembic_cfg - - -def delete_nodes_and_connections_sqla(pks_to_delete): - """ - Delete all nodes corresponding to pks in the input. - :param pks_to_delete: A list, tuple or set of pks that should be deleted. - """ - from aiida.backends import sqlalchemy as sa - from aiida.backends.sqlalchemy.models.node import DbNode, DbLink - from aiida.backends.sqlalchemy.models.group import table_groups_nodes - from aiida.manage.manager import get_manager - - backend = get_manager().get_backend() - - with backend.transaction() as session: - # I am first making a statement to delete the membership of these nodes to groups. - # Since table_groups_nodes is a sqlalchemy.schema.Table, I am using expression language to compile - # a stmt to be executed by the session. It works, but it's not nice that two different ways are used! - # Can this be changed? - stmt = table_groups_nodes.delete().where(table_groups_nodes.c.dbnode_id.in_(list(pks_to_delete))) - session.execute(stmt) - # First delete links, then the Nodes, since we are not cascading deletions. - # Here I delete the links coming out of the nodes marked for deletion. - session.query(DbLink).filter(DbLink.input_id.in_(list(pks_to_delete))).delete(synchronize_session='fetch') - # Here I delete the links pointing to the nodes marked for deletion. - session.query(DbLink).filter(DbLink.output_id.in_(list(pks_to_delete))).delete(synchronize_session='fetch') - # Now I am deleting the nodes - session.query(DbNode).filter(DbNode.id.in_(list(pks_to_delete))).delete(synchronize_session='fetch') +""" + ) + return pg_tc.substitute( + links_table_name=links_table_name, + links_table_input_field=links_table_input_field, + links_table_output_field=links_table_output_field, + closure_table_name=closure_table_name, + closure_table_parent_field=closure_table_parent_field, + closure_table_child_field=closure_table_child_field + ) diff --git a/aiida/backends/utils.py b/aiida/backends/utils.py index 3908aa06bb..b2ce3ec80b 100644 --- a/aiida/backends/utils.py +++ b/aiida/backends/utils.py @@ -11,139 +11,11 @@ from __future__ import print_function from __future__ import absolute_import -import abc -import collections -import six - from aiida.backends import BACKEND_SQLA, BACKEND_DJANGO -from aiida.common.exceptions import ConfigurationError, ValidationError from aiida.manage import configuration AIIDA_ATTRIBUTE_SEP = '.' -SCHEMA_GENERATION_KEY = 'schema_generation' -SCHEMA_GENERATION_VALUE = '1' - - -Setting = collections.namedtuple('Setting', ['key', 'value', 'description', 'time']) - - -class SettingsManager(object): - """Class to get, set and delete settings from the `DbSettings` table.""" - - @abc.abstractmethod - def get(self, key): - """Return the setting with the given key. - - :param key: the key identifying the setting - :return: Setting - :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist - """ - - @abc.abstractmethod - def set(self, key, value, description=None): - """Return the settings with the given key. - - :param key: the key identifying the setting - :param value: the value for the setting - :param description: optional setting description - """ - - @abc.abstractmethod - def delete(self, key): - """Delete the setting with the given key. - - :param key: the key identifying the setting - :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist - """ - - -def get_settings_manager(): - if configuration.PROFILE.database_backend == BACKEND_DJANGO: - from aiida.backends.djsite.utils import DjangoSettingsManager - manager = DjangoSettingsManager() - elif configuration.PROFILE.database_backend == BACKEND_SQLA: - from aiida.backends.sqlalchemy.utils import SqlaSettingsManager - manager = SqlaSettingsManager() - else: - raise Exception('unknown backend type `{}`'.format(configuration.PROFILE.database_backend)) - - return manager - - -def validate_schema_generation(): - """Verify that the database schema generation is compatible with the current code schema generation.""" - from aiida.common.exceptions import ConfigurationError, NotExistent - try: - schema_generation_database = get_db_schema_generation().value - except NotExistent: - schema_generation_database = '1' - - if schema_generation_database is None: - schema_generation_database = '1' - - if schema_generation_database != SCHEMA_GENERATION_VALUE: - raise ConfigurationError( - 'The schema generation of your database {} is newer than that of the code `{}` and is incompatible.'.format( - schema_generation_database, SCHEMA_GENERATION_VALUE - ) - ) - - -def get_db_schema_generation(): - """Get the schema generation of the current database.""" - from aiida.backends.utils import get_settings_manager - manager = get_settings_manager() - return manager.get(SCHEMA_GENERATION_KEY) - - -def validate_attribute_key(key): - """ - Validate the key string to check if it is valid (e.g., if it does not - contain the separator symbol.). - - :return: None if the key is valid - :raise aiida.common.ValidationError: if the key is not valid - """ - if not isinstance(key, six.string_types): - raise ValidationError('The key must be a string.') - if not key: - raise ValidationError('The key cannot be an empty string.') - if AIIDA_ATTRIBUTE_SEP in key: - raise ValidationError("The separator symbol '{}' cannot be present " - 'in the key of attributes, extras, etc.'.format(AIIDA_ATTRIBUTE_SEP)) - - -def load_dbenv(profile=None, *args, **kwargs): - if configuration.PROFILE.database_backend == BACKEND_SQLA: - # Maybe schema version should be also checked for SQLAlchemy version. - from aiida.backends.sqlalchemy.utils \ - import load_dbenv as load_dbenv_sqlalchemy - to_return = load_dbenv_sqlalchemy(profile=profile, *args, **kwargs) - elif configuration.PROFILE.database_backend == BACKEND_DJANGO: - from aiida.backends.djsite.utils import load_dbenv as load_dbenv_django - to_return = load_dbenv_django(profile=profile, *args, **kwargs) - else: - raise ConfigurationError('Invalid configuration.PROFILE.database_backend: {}'.format( - configuration.PROFILE.database_backend)) - - return to_return - - -def _load_dbenv_noschemacheck(profile=None, *args, **kwargs): - if configuration.PROFILE.database_backend == BACKEND_SQLA: - # Maybe schema version should be also checked for SQLAlchemy version. - from aiida.backends.sqlalchemy.utils \ - import _load_dbenv_noschemacheck as _load_dbenv_noschemacheck_sqlalchemy - to_return = _load_dbenv_noschemacheck_sqlalchemy(profile=profile, *args, **kwargs) - elif configuration.PROFILE.database_backend == BACKEND_DJANGO: - from aiida.backends.djsite.utils import _load_dbenv_noschemacheck as _load_dbenv_noschemacheck_django - to_return = _load_dbenv_noschemacheck_django(profile=profile, *args, **kwargs) - else: - raise ConfigurationError('Invalid configuration.PROFILE.database_backend: {}'.format(configuration.PROFILE.database_backend)) - - return to_return - def delete_nodes_and_connections(pks): if configuration.PROFILE.database_backend == BACKEND_DJANGO: diff --git a/aiida/cmdline/commands/cmd_database.py b/aiida/cmdline/commands/cmd_database.py index e75613fc39..b0771cba75 100644 --- a/aiida/cmdline/commands/cmd_database.py +++ b/aiida/cmdline/commands/cmd_database.py @@ -14,6 +14,7 @@ import click +from aiida.common import exceptions from aiida.cmdline.commands.cmd_verdi import verdi from aiida.cmdline.params import options from aiida.cmdline.utils import decorators, echo @@ -41,7 +42,10 @@ def database_migrate(force): backend = manager._load_backend(schema_check=False) # pylint: disable=protected-access if force: - backend.migrate() + try: + backend.migrate() + except exceptions.ConfigurationError as exception: + echo.echo_critical(str(exception)) return echo.echo_warning('Migrating your database might take a while and is not reversible.') @@ -67,8 +71,12 @@ def database_migrate(force): echo.echo('\n') echo.echo_critical('Migration aborted, the data has not been affected.') else: - backend.migrate() - echo.echo_success('migration completed') + try: + backend.migrate() + except exceptions.ConfigurationError as exception: + echo.echo_critical(str(exception)) + else: + echo.echo_success('migration completed') @verdi_database.group('integrity') diff --git a/aiida/common/exceptions.py b/aiida/common/exceptions.py index 6e6074e658..47b0f5ffa5 100644 --- a/aiida/common/exceptions.py +++ b/aiida/common/exceptions.py @@ -18,9 +18,9 @@ 'EntryPointError', 'MissingEntryPointError', 'MultipleEntryPointError', 'LoadingEntryPointError', 'InvalidEntryPointTypeError', 'InvalidOperation', 'ParsingError', 'InternalError', 'PluginInternalError', 'ValidationError', 'ConfigurationError', 'ProfileConfigurationError', 'MissingConfigurationError', - 'ConfigurationVersionError', 'DbContentError', 'InputValidationError', 'FeatureNotAvailable', 'FeatureDisabled', - 'LicensingException', 'TestsNotAllowedError', 'UnsupportedSpeciesError', 'TransportTaskException', - 'OutputParsingError' + 'ConfigurationVersionError', 'IncompatibleDatabaseSchema', 'DbContentError', 'InputValidationError', + 'FeatureNotAvailable', 'FeatureDisabled', 'LicensingException', 'TestsNotAllowedError', 'UnsupportedSpeciesError', + 'TransportTaskException', 'OutputParsingError' ) @@ -173,6 +173,10 @@ class ConfigurationVersionError(ConfigurationError): """ +class IncompatibleDatabaseSchema(ConfigurationError): + """Raised when the database schema is incompatible with that of the code.""" + + class DbContentError(AiidaException): """ Raised when the content of the DB is not valid. diff --git a/aiida/engine/utils.py b/aiida/engine/utils.py index 311b3f47ee..39b55ff47e 100644 --- a/aiida/engine/utils.py +++ b/aiida/engine/utils.py @@ -242,9 +242,9 @@ def set_process_state_change_timestamp(process): :param process: the Process instance that changed its state """ - from aiida.backends.utils import get_settings_manager from aiida.common import timezone from aiida.common.exceptions import UniquenessError + from aiida.manage.manager import get_manager # pylint: disable=cyclic-import from aiida.orm import ProcessNode, CalculationNode, WorkflowNode if isinstance(process.node, CalculationNode): @@ -262,8 +262,8 @@ def set_process_state_change_timestamp(process): value = timezone.datetime_to_isoformat(timezone.now()) try: - manager = get_settings_manager() - manager.set(key, value, description) + manager = get_manager() + manager.get_backend_manager().get_settings_manager().set(key, value, description) except UniquenessError as exception: process.logger.debug('could not update the {} setting because of a UniquenessError: {}'.format(key, exception)) @@ -278,11 +278,11 @@ def get_process_state_change_timestamp(process_type=None): known process types will be returned. :return: a timestamp or None """ - from aiida.backends.utils import get_settings_manager from aiida.common import timezone from aiida.common.exceptions import NotExistent + from aiida.manage.manager import get_manager # pylint: disable=cyclic-import - manager = get_settings_manager() + manager = get_manager().get_backend_manager().get_settings_manager() valid_process_types = ['calculation', 'work'] if process_type is not None and process_type not in valid_process_types: diff --git a/aiida/manage/manager.py b/aiida/manage/manager.py index d7d1067c42..6f247e6c6d 100644 --- a/aiida/manage/manager.py +++ b/aiida/manage/manager.py @@ -48,22 +48,8 @@ def get_profile(): def unload_backend(self): """Unload the current backend and its corresponding database environment.""" - from aiida.backends import BACKEND_DJANGO, BACKEND_SQLA - from aiida.common import ConfigurationError - - profile = self.get_profile() - backend = profile.database_backend - - if backend == BACKEND_DJANGO: - from aiida.backends.djsite.utils import unload_dbenv - unload_backend = unload_dbenv - elif backend == BACKEND_SQLA: - from aiida.backends.sqlalchemy.utils import unload_dbenv - unload_backend = unload_dbenv - else: - raise ConfigurationError('Invalid backend type {} in profile: {}'.format(backend, profile.name)) - - unload_backend() + manager = self.get_backend_manager() + manager.reset_backend_environment() self._backend = None def _load_backend(self, schema_check=True): @@ -91,22 +77,14 @@ def _load_backend(self, schema_check=True): if configuration.BACKEND_UUID is not None and configuration.BACKEND_UUID != profile.uuid: raise InvalidOperation('cannot load backend because backend of another profile is already loaded') - backend_type = profile.database_backend - - if backend_type == BACKEND_DJANGO: - from aiida.backends.djsite.utils import load_dbenv, _load_dbenv_noschemacheck - load_backend = load_dbenv if schema_check else _load_dbenv_noschemacheck - elif backend_type == BACKEND_SQLA: - from aiida.backends.sqlalchemy.utils import load_dbenv, _load_dbenv_noschemacheck - load_backend = load_dbenv if schema_check else _load_dbenv_noschemacheck - else: - raise ConfigurationError('Invalid backend type {} in profile: {}'.format(backend_type, profile.name)) - # Do NOT reload the backend environment if already loaded, simply reload the backend instance after if configuration.BACKEND_UUID is None: - load_backend(profile) + manager = self.get_backend_manager() + manager.load_backend_environment(profile, validate_schema=schema_check) configuration.BACKEND_UUID = profile.uuid + backend_type = profile.database_backend + # Can only import the backend classes after the backend has been loaded if backend_type == BACKEND_DJANGO: from aiida.orm.implementation.django.backend import DjangoBackend @@ -129,9 +107,23 @@ def backend_loaded(self): """ return self._backend is not None - def get_backend(self): + def get_backend_manager(self): + """Return the database backend manager. + + .. note:: this is not the actual backend, but a manager class that is necessary for database operations that + go around the actual ORM. For example when the schema version has not yet been validated. + + :return: the database backend manager + :rtype: :class:`aiida.backend.manager.BackendManager` """ - Get the database backend + if self._backend_manager is None: + from aiida.backends import get_backend_manager + self._backend_manager = get_backend_manager(self.get_profile().database_backend) + + return self._backend_manager + + def get_backend(self): + """Return the database backend :return: the database backend :rtype: :class:`aiida.orm.implementation.Backend` @@ -142,8 +134,7 @@ def get_backend(self): return self._backend def get_persister(self): - """ - Get the persister + """Return the persister :return: the current persister instance :rtype: :class:`plumpy.Persister` @@ -156,8 +147,7 @@ def get_persister(self): return self._persister def get_communicator(self): - """ - Get the communicator + """Return the communicator :return: a global communicator instance :rtype: :class:`kiwipy.Communicator` @@ -168,8 +158,7 @@ def get_communicator(self): return self._communicator def create_communicator(self, task_prefetch_count=None, with_orm=True): - """ - Create a Communicator + """Create a Communicator :param task_prefetch_count: optional specify how many tasks this communicator take simultaneously :param with_orm: if True, use ORM (de)serializers. If false, use json. @@ -218,8 +207,7 @@ def create_communicator(self, task_prefetch_count=None, with_orm=True): ) def get_daemon_client(self): - """ - Return the daemon client for the current profile. + """Return the daemon client for the current profile. :return: the daemon client :rtype: :class:`aiida.daemon.client.DaemonClient` @@ -234,8 +222,7 @@ def get_daemon_client(self): return self._daemon_client def get_process_controller(self): - """ - Get a process controller + """Return the process controller :return: the process controller instance :rtype: :class:`plumpy.RemoteProcessThreadController` @@ -247,8 +234,7 @@ def get_process_controller(self): return self._process_controller def get_runner(self): - """ - Get a runner that is based on the current profile settings and can be used globally by the code. + """Return a runner that is based on the current profile settings and can be used globally by the code. :return: the global runner :rtype: :class:`aiida.engine.runners.Runner` @@ -259,8 +245,7 @@ def get_runner(self): return self._runner def set_runner(self, new_runner): - """ - Set the currently used runner + """Set the currently used runner :param new_runner: the new runner to use :type new_runner: :class:`aiida.engine.runners.Runner` @@ -271,8 +256,7 @@ def set_runner(self, new_runner): self._runner = new_runner def create_runner(self, with_persistence=True, **kwargs): - """ - Create a new runner + """Create and return a new runner :param with_persistence: create a runner with persistence enabled :type with_persistence: bool @@ -299,8 +283,9 @@ def create_runner(self, with_persistence=True, **kwargs): return runners.Runner(**settings) def create_daemon_runner(self, loop=None): - """ - Create a new daemon runner. This is used by workers when the daemon is running and in testing. + """Create and return a new daemon runner. + + This is used by workers when the daemon is running and in testing. :param loop: the (optional) tornado event loop to use :type loop: :class:`tornado.ioloop.IOLoop` @@ -329,15 +314,14 @@ def callback(*args, **kwargs): return runner def close(self): - """ - Reset the global settings entirely and release any global objects - """ + """Reset the global settings entirely and release any global objects.""" if self._communicator is not None: self._communicator.stop() if self._runner is not None: self._runner.stop() self._backend = None + self._backend_manager = None self._config = None self._profile = None self._communicator = None @@ -349,6 +333,7 @@ def close(self): def __init__(self): super(Manager, self).__init__() self._backend = None # type: aiida.orm.implementation.Backend + self._backend_manager = None # type: aiida.backend.manager.BackendManager self._config = None # type: aiida.manage.configuration.config.Config self._daemon_client = None # type: aiida.daemon.client.DaemonClient self._profile = None # type: aiida.manage.configuration.profile.Profile diff --git a/aiida/orm/implementation/backends.py b/aiida/orm/implementation/backends.py index 2d7af87133..71d9971729 100644 --- a/aiida/orm/implementation/backends.py +++ b/aiida/orm/implementation/backends.py @@ -27,7 +27,7 @@ class Backend(object): @abc.abstractmethod def migrate(self): - """Migrate the database to the latest schema version.""" + """Migrate the database to the latest schema generation or version.""" @abc.abstractproperty def authinfos(self): diff --git a/aiida/orm/implementation/django/backend.py b/aiida/orm/implementation/django/backend.py index be0361d7c9..77a36fc79a 100644 --- a/aiida/orm/implementation/django/backend.py +++ b/aiida/orm/implementation/django/backend.py @@ -18,7 +18,7 @@ from django.db import models, transaction from aiida.backends.djsite.queries import DjangoQueryManager -from aiida.backends.djsite.utils import migrate_database +from aiida.backends.djsite.manager import DjangoBackendManager from ..sql import SqlBackend from . import authinfos @@ -46,11 +46,11 @@ def __init__(self): self._logs = logs.DjangoLogCollection(self) self._nodes = nodes.DjangoNodeCollection(self) self._query_manager = DjangoQueryManager(self) + self._backend_manager = DjangoBackendManager() self._users = users.DjangoUserCollection(self) - @staticmethod - def migrate(): - migrate_database() + def migrate(self): + self._backend_manager.migrate() @property def authinfos(self): diff --git a/aiida/orm/implementation/sqlalchemy/backend.py b/aiida/orm/implementation/sqlalchemy/backend.py index 9587547dd8..0dadebc1d6 100644 --- a/aiida/orm/implementation/sqlalchemy/backend.py +++ b/aiida/orm/implementation/sqlalchemy/backend.py @@ -17,7 +17,7 @@ from aiida.backends.sqlalchemy import get_scoped_session from aiida.backends.sqlalchemy.models import base from aiida.backends.sqlalchemy.queries import SqlaQueryManager -from aiida.backends.sqlalchemy.utils import migrate_database +from aiida.backends.sqlalchemy.manager import SqlaBackendManager from ..sql import SqlBackend from . import authinfos @@ -45,11 +45,11 @@ def __init__(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) - @staticmethod - def migrate(): - migrate_database() + def migrate(self): + self._schema_manager.migrate() @property def authinfos(self): @@ -124,21 +124,27 @@ def cursor(self): finally: self.get_connection().close() - @staticmethod - def execute_raw(query): + 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 """ - session = get_scoped_session() - result = session.execute(query) - return result.fetchall() + 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 + """Get the SQLA database connection :return: the SQLA database connection """ diff --git a/aiida/plugins/factories.py b/aiida/plugins/factories.py index 02eadd8c29..447135128f 100644 --- a/aiida/plugins/factories.py +++ b/aiida/plugins/factories.py @@ -7,7 +7,7 @@ # For further information on the license, see the LICENSE.txt file # # For further information please visit http://www.aiida.net # ########################################################################### -# pylint: disable=invalid-name,inconsistent-return-statements +# pylint: disable=invalid-name,inconsistent-return-statements,cyclic-import """Definition of factories to load classes from the various plugin groups.""" from __future__ import division from __future__ import print_function diff --git a/docs/source/nitpick-exceptions b/docs/source/nitpick-exceptions index 2b2654824c..aecf291b63 100644 --- a/docs/source/nitpick-exceptions +++ b/docs/source/nitpick-exceptions @@ -272,3 +272,6 @@ py:class aldjemy.orm.DbAuthInfo py:class aldjemy.orm.DbComment py:class aldjemy.orm.DbLog py:class aldjemy.orm.DbSetting + +# Alembic +py:class alembic.config.Config \ No newline at end of file