From 952a5e9fc0ded86692f5df5986db190a9ce5a438 Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Fri, 26 Jul 2019 19:02:32 +0200 Subject: [PATCH] Refactor the code dealing with database into `BackendManager` The code that acts directly with the database, outside of the ORM, for example to check the database schema generation and version, or to get data from the `DbSettings` table has been refactored. The ORM requires the database environment to be loaded, but to do that the database schema needs to be verified which also requires accessing the database. This functionality was implemented mostly in free-standing functions and separately for both backends. This is reorganized into a `BackendManager` abstract class and implemented for Django and SqlAlchemy. This centralizes the functionality required to verify database schema and perform the migrations thereof. --- .pre-commit-config.yaml | 4 - aiida/backends/__init__.py | 17 + .../backends/djsite/db/migrations/__init__.py | 54 ++- ...ns_0037_attributes_extras_settings_json.py | 14 +- aiida/backends/djsite/manage.py | 14 +- aiida/backends/djsite/manager.py | 190 +++++++++++ aiida/backends/djsite/utils.py | 200 +---------- aiida/backends/manager.py | 313 ++++++++++++++++++ aiida/backends/sqlalchemy/__init__.py | 7 +- aiida/backends/sqlalchemy/manage.py | 12 +- aiida/backends/sqlalchemy/manager.py | 210 ++++++++++++ aiida/backends/sqlalchemy/models/settings.py | 2 +- aiida/backends/sqlalchemy/queries.py | 2 - .../sqlalchemy/tests/test_migrations.py | 21 +- aiida/backends/sqlalchemy/utils.py | 307 +++-------------- aiida/backends/utils.py | 128 ------- aiida/cmdline/commands/cmd_database.py | 14 +- aiida/common/exceptions.py | 10 +- aiida/engine/utils.py | 10 +- aiida/manage/manager.py | 87 ++--- aiida/orm/implementation/backends.py | 2 +- aiida/orm/implementation/django/backend.py | 8 +- .../orm/implementation/sqlalchemy/backend.py | 28 +- aiida/plugins/factories.py | 2 +- docs/source/nitpick-exceptions | 3 + 25 files changed, 945 insertions(+), 714 deletions(-) create mode 100644 aiida/backends/djsite/manager.py create mode 100644 aiida/backends/manager.py create mode 100644 aiida/backends/sqlalchemy/manager.py 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