🧪 TESTS: SQLA Migrations -> pytest

aiida/backends/djsite/manager.py

@@ -63,10 +63,9 @@ def is_database_schema_ahead(self):
         """
         # 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())
+        return StrictVersion(self.get_schema_version_backend()) > StrictVersion(self.get_schema_version_head())
 
-    def get_schema_version_code(self):
-        """Return the code schema version."""
+    def get_schema_version_head(self):
         from .db.models import SCHEMA_VERSION
         return SCHEMA_VERSION
 
@@ -101,11 +100,7 @@ def get_schema_generation_database(self):
             except (IndexError, ValueError, TypeError):
                 return '1'
 
-    def get_schema_version_database(self):
-        """Return the database schema version.
-
-        :return: `distutils.version.StrictVersion` with schema version of the database
-        """
+    def get_schema_version_backend(self):
         from django.db.utils import ProgrammingError
 
         from aiida.manage.manager import get_manager

aiida/backends/manager.py

@@ -145,7 +145,7 @@ def migrate(self):
             self.validate_schema_generation_for_migration()
             self._migrate_database_generation()
 
-        if self.get_schema_version_code() != self.get_schema_version_database():
+        if self.get_schema_version_head() != self.get_schema_version_backend():
             self.validate_schema_version_for_migration()
             self._migrate_database_version()
 
@@ -160,7 +160,7 @@ def _migrate_database_generation(self):
         logic that is required.
         """
         self.set_schema_generation_database(SCHEMA_GENERATION_VALUE)
-        self.set_schema_version_database(self.get_schema_version_code())
+        self.set_schema_version_database(self.get_schema_version_head())
 
     def _migrate_database_version(self):
         """Migrate the database to the current schema version.
@@ -179,8 +179,8 @@ def is_database_schema_ahead(self):
         """
 
     @abc.abstractmethod
-    def get_schema_version_code(self):
-        """Return the code schema version."""
+    def get_schema_version_head(self):
+        """Return the latest schema version."""
 
     @abc.abstractmethod
     def get_schema_version_reset(self, schema_generation_code):
@@ -191,11 +191,8 @@ def get_schema_version_reset(self, schema_generation_code):
         """
 
     @abc.abstractmethod
-    def get_schema_version_database(self):
-        """Return the database schema version.
-
-        :return: `distutils.version.LooseVersion` with schema version of the database
-        """
+    def get_schema_version_backend(self):
+        """Return the current schema version."""
 
     @abc.abstractmethod
     def set_schema_version_database(self, version):
@@ -256,7 +253,7 @@ def validate_schema_generation_for_migration(self):
         """
         schema_generation_code = SCHEMA_GENERATION_VALUE
         schema_generation_database = self.get_schema_generation_database()
-        schema_version_database = self.get_schema_version_database()
+        schema_version_database = self.get_schema_version_backend()
         schema_version_reset = self.get_schema_version_reset(schema_generation_code)
         schema_generation_reset, aiida_core_version_reset = SCHEMA_GENERATION_RESET[schema_generation_code]
 
@@ -287,8 +284,8 @@ def validate_schema_version_for_migration(self):
 
         :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()
+        schema_version_code = self.get_schema_version_head()
+        schema_version_database = self.get_schema_version_backend()
 
         if self.is_database_schema_ahead():
             # Database is newer than the code so a downgrade would be necessary but this is not supported.
@@ -322,8 +319,8 @@ def validate_schema_version(self, profile):
         :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()
+        schema_version_code = self.get_schema_version_head()
+        schema_version_database = self.get_schema_version_backend()
 
         if schema_version_database != schema_version_code:
             raise exceptions.IncompatibleDatabaseSchema(

aiida/backends/sqlalchemy/manager.py

@@ -11,6 +11,7 @@
 import contextlib
 import os
 
+from alembic.command import downgrade, upgrade
 import sqlalchemy
 from sqlalchemy.orm.exc import NoResultFound
 
@@ -94,6 +95,14 @@ def reset_backend_environment(self):
         from . import reset_session
         reset_session()
 
+    def list_schema_versions(self):
+        """List all available schema versions (oldest to latest).
+
+        :return: list of strings with schema versions
+        """
+        with self.alembic_script() as script:
+            return list(reversed([entry.revision for entry in script.walk_revisions()]))
+
     def is_database_schema_ahead(self):
         """Determine whether the database schema version is ahead of the code schema version.
 
@@ -102,10 +111,9 @@ def is_database_schema_ahead(self):
         :return: boolean, True if the database schema version is ahead of the code schema version.
         """
         with self.alembic_script() as script:
-            return self.get_schema_version_database() not in [entry.revision for entry in script.walk_revisions()]
+            return self.get_schema_version_backend() not in [entry.revision for entry in script.walk_revisions()]
 
-    def get_schema_version_code(self):
-        """Return the code schema version."""
+    def get_schema_version_head(self):
         with self.alembic_script() as script:
             return script.get_current_head()
 
@@ -117,11 +125,7 @@ def get_schema_version_reset(self, schema_generation_code):
         """
         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_schema_version_backend(self):
         with self.migration_context() as context:
             return context.get_current_revision()
 
@@ -133,13 +137,26 @@ def set_schema_version_database(self, version):
         with self.migration_context() as context:
             return context.stamp(context.script, 'head')
 
-    def _migrate_database_version(self):
-        """Migrate the database to the current schema version."""
-        super()._migrate_database_version()
-        from alembic.command import upgrade
+    def migrate_up(self, version: str):
+        """Migrate the database up to a specific version.
 
+        :param version: string with schema version to migrate to
+        """
+        with self.alembic_config() as config:
+            upgrade(config, version)
+
+    def migrate_down(self, version: str):
+        """Migrate the database down to a specific version.
+
+        :param version: string with schema version to migrate to
+        """
         with self.alembic_config() as config:
-            upgrade(config, 'head')
+            downgrade(config, version)
+
+    def _migrate_database_version(self):
+        """Migrate the database to the latest schema version."""
+        super()._migrate_database_version()
+        self.migrate_up('head')
 
 
 class SqlaSettingsManager(SettingsManager):

aiida/cmdline/commands/cmd_database.py

@@ -38,7 +38,7 @@ def database_version():
     echo.echo('Generation: ', bold=True, nl=False)
     echo.echo(backend_manager.get_schema_generation_database())
     echo.echo('Revision:   ', bold=True, nl=False)
-    echo.echo(backend_manager.get_schema_version_database())
+    echo.echo(backend_manager.get_schema_version_backend())
 
 
 @verdi_database.command('migrate')

aiida/manage/tests/main.py

@@ -117,8 +117,8 @@ def use_profile(self, profile_name):
     def has_profile_open(self):
         return self._manager and self._manager.has_profile_open()
 
-    def reset_db(self):
-        return self._manager.reset_db()
+    def reset_db(self, init_db=True):
+        return self._manager.reset_db(init_db=init_db)
 
     def destroy_all(self):
         if self._manager:
@@ -166,10 +166,11 @@ def _select_db_test_case(self, backend):
             self._test_case = SqlAlchemyTests()
             self._test_case.test_session = get_scoped_session()
 
-    def reset_db(self):
+    def reset_db(self, init_db=True):
         self._test_case.clean_db()  # will drop all users
         manager.reset_manager()
-        self.init_db()
+        if init_db:
+            self.init_db()
 
     def init_db(self):
         """Initialise the database state for running of tests.

tests/backends/aiida_sqlalchemy/migrations/conftest.py

@@ -0,0 +1,90 @@
+# -*- 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               #
+###########################################################################
+"""Tests for the migration engine (Alembic) as well as for the AiiDA migrations for SQLAlchemy."""
+from contextlib import contextmanager
+from typing import Iterator
+
+import pytest
+from sqlalchemy.orm import Session
+
+from aiida.backends.sqlalchemy.manager import SqlaBackendManager
+
+
+class Migrator:
+    """A class to yield from the ``perform_migrations`` fixture."""
+
+    def __init__(self, backend, manager: SqlaBackendManager) -> None:
+        self.backend = backend
+        self._manager = manager
+
+    def migrate_up(self, revision: str) -> None:
+        """Migrate up to a given revision."""
+        self._manager.migrate_up(revision)
+        assert self._manager.get_schema_version_backend() == revision
+
+    def migrate_down(self, revision: str) -> None:
+        """Migrate down to a given revision."""
+        self._manager.migrate_down(revision)
+        assert self._manager.get_schema_version_backend() == revision
+
+    def get_current_table(self, table_name):
+        """
+        Return a Model instantiated at the correct migration.
+        Note that this is obtained by inspecting the database and not
+        by looking into the models file.
+        So, special methods possibly defined in the models files/classes are not present.
+
+        For instance, you can do::
+
+          DbGroup = self.get_current_table('db_dbgroup')
+
+        :param table_name: the name of the table.
+        """
+        from alembic.migration import MigrationContext  # pylint: disable=import-error
+        from sqlalchemy.ext.automap import automap_base  # pylint: disable=import-error,no-name-in-module
+
+        with self.backend.get_session().bind.begin() as connection:
+            context = MigrationContext.configure(connection)
+            bind = context.bind
+
+            base = automap_base()
+            # reflect the tables
+            base.prepare(bind.engine, reflect=True)
+
+            return getattr(base.classes, table_name)
+
+    @contextmanager
+    def session(self) -> Iterator[Session]:
+        """A context manager for a new session."""
+        with self.backend.get_session().bind.begin() as connection:
+            session = Session(connection.engine, future=True)
+            try:
+                yield session
+            except Exception:
+                session.rollback()
+                raise
+            finally:
+                session.close()
+
+
+@pytest.fixture()
+def perform_migrations(aiida_profile, backend, request):
+    """A fixture to setup the database for migration tests"""
+    # note downgrading to 1830c8430131 requires adding columns to `DbUser` and hangs if a user is present
+    aiida_profile.reset_db(init_db=False)
+    migrator = Migrator(backend, SqlaBackendManager())
+    marker = request.node.get_closest_marker('migrate_down')
+    if marker is not None:
+        assert marker.args, 'No version given'
+        migrator.migrate_down(marker.args[0])
+    yield migrator
+    # ensure that the database is migrated back up to the latest version, once finished
+    aiida_profile.reset_db(init_db=False)
+    SqlaBackendManager().migrate_up('head')

tests/backends/aiida_sqlalchemy/migrations/test_all_basic.py

@@ -0,0 +1,35 @@
+# -*- 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               #
+###########################################################################
+"""Basic tests for all migratios"""
+import pytest
+
+
+@pytest.mark.usefixtures('perform_migrations')
+def test_all_empty_migrations():
+    """Test migrating down to a particular version, then back up, using an empty database.
+
+    Note, migrating down+up with 59edaf8a8b79_adding_indexes_and_constraints_to_the_.py raises::
+
+        sqlalchemy.exc.ProgrammingError:
+        (psycopg2.errors.DuplicateTable) relation "db_dbgroup_dbnodes_dbgroup_id_dbnode_id_key" already exists
+
+    So we only run for all versions later than this.
+    """
+    from aiida.backends.sqlalchemy.manager import SqlaBackendManager
+    migrator = SqlaBackendManager()
+    all_versions = migrator.list_schema_versions()
+    first_index = all_versions.index('a514d673c163') + 1
+    # ideally we would pytest parametrize this, but then we would need to call list_schema_versions on module load
+    for version in all_versions[first_index:]:
+        print(version)
+        migrator.migrate_down(version)
+        assert migrator.get_schema_version_backend() == version
+        migrator.migrate_up('head')
+        assert migrator.get_schema_version_backend() == migrator.get_schema_version_head()