Refactor the code dealing with database into BackendManager

.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|

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

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):
         """

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(

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

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