Source code for

# Copyright (c), The AiiDA team. All rights reserved.                     #
# This file is part of the AiiDA code.                                    #
#                                                                         #
# The code is hosted on GitHub at #
# For further information on the license, see the LICENSE.txt file        #
# For further information please visit               #
"""Schema validation and migration utilities.

This code interacts directly with the database, outside of the ORM,
taking a `Profile` as input for the connection configuration.

.. important:: This code should only be accessed via the storage backend class, not directly!

from __future__ import annotations

import contextlib
import pathlib
from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional

from alembic.command import downgrade, upgrade
from alembic.config import Config
from alembic.runtime.environment import EnvironmentContext
from alembic.runtime.migration import MigrationContext, MigrationInfo
from alembic.script import ScriptDirectory
from sqlalchemy import MetaData, String, column, desc, insert, inspect, select, table
from sqlalchemy.exc import OperationalError, ProgrammingError
from sqlalchemy.ext.automap import automap_base
from sqlalchemy.orm import Session

from aiida.common import exceptions
from aiida.manage.configuration.profile import Profile
from import MIGRATE_LOGGER
from import DbSetting
from import create_sqlalchemy_engine

    from disk_objectstore import Container

Database schema is using the legacy Django schema.
To migrate the database schema version to the current one, run the following command:

    verdi -p {profile_name} storage migrate

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} storage migrate
"""  # noqa: E501

ALEMBIC_REL_PATH = 'migrations'

REPOSITORY_UUID_KEY = 'repository|uuid'

[docs] class PsqlDosMigrator: """Class for validating and migrating `psql_dos` storage instances. .. important:: This class should only be accessed via the storage backend class (apart from for test purposes) """ alembic_version_tbl_name = 'alembic_version' django_version_table = table( 'django_migrations', column('id'), column('app', String(255)), column('name', String(255)), column('applied') )
[docs] def __init__(self, profile: Profile) -> None: self.profile = profile self._engine = create_sqlalchemy_engine(self.profile.storage_config) self._connection = None
[docs] def close(self) -> None: """Close the connection if it was opened and dispose of the engine.""" if self._connection: self._connection.close() self._connection = None if self._engine: self._engine.dispose() self._engine = None
@property def connection(self): """Return the connection to the database. Will automatically create the engine and open an connection if not already opened in a previous call. :return: Open connection to the database. :raises: :class:`aiida.common.exceptions.UnreachableStorage` if connecting to the database fails. """ if self._connection is None: try: self._connection = self._engine.connect() except OperationalError as exception: raise exceptions.UnreachableStorage(f'Could not connect to database: {exception}') from exception return self._connection
[docs] @classmethod def get_schema_versions(cls) -> Dict[str, str]: """Return all available schema versions (oldest to latest). :return: schema version -> description """ return {entry.revision: entry.doc for entry in reversed(list(cls._alembic_script().walk_revisions()))}
[docs] @classmethod def get_schema_version_head(cls) -> str: """Return the head schema version for this storage, i.e. the latest schema this storage can be migrated to.""" return cls._alembic_script().revision_map.get_current_head('main')
[docs] def get_schema_version_profile(self, check_legacy=False) -> Optional[str]: """Return the schema version of the backend instance for this profile. Note, the version will be None if the database is empty or is a legacy django database. """ with self._migration_context() as context: version = context.get_current_revision() if version is None and check_legacy: stmt = select( == 'db') stmt = stmt.order_by(desc( try: return self.connection.execute(stmt).scalar() except (OperationalError, ProgrammingError): self.connection.rollback() return version
[docs] def validate_storage(self) -> None: """Validate that the storage for this profile 1. That the database schema is at the head version, i.e. is compatible with the code API. 2. That the repository ID is equal to the UUID set in the database :raises: :class:`aiida.common.exceptions.UnreachableStorage` if the storage cannot be connected to :raises: :class:`aiida.common.exceptions.IncompatibleStorageSchema` if the storage is not compatible with the code API. :raises: :class:`aiida.common.exceptions.CorruptStorage` if the repository ID is not equal to the UUID set in thedatabase. """ # check there is an alembic_version table from which to get the schema version if not inspect(self.connection).has_table(self.alembic_version_tbl_name): # if not present, it might be that this is a legacy django database if inspect(self.connection).has_table( raise exceptions.IncompatibleStorageSchema( TEMPLATE_LEGACY_DJANGO_SCHEMA.format( ) raise exceptions.IncompatibleStorageSchema('The database has no known version.') # now we can check that the alembic version is the latest schema_version_code = self.get_schema_version_head() schema_version_database = self.get_schema_version_profile(check_legacy=False) if schema_version_database != schema_version_code: raise exceptions.IncompatibleStorageSchema( TEMPLATE_INVALID_SCHEMA_VERSION.format( schema_version_database=schema_version_database, schema_version_code=schema_version_code,, ) ) # finally, we check that the ID set within the disk-objectstore is equal to the one saved in the database, # i.e. this container is indeed the one associated with the db repository_uuid = self.get_repository_uuid() stmt = select(DbSetting.val).where(DbSetting.key == REPOSITORY_UUID_KEY) database_repository_uuid = self.connection.execute(stmt).scalar_one_or_none() if database_repository_uuid is None: raise exceptions.CorruptStorage('The database has no repository UUID set.') if database_repository_uuid != repository_uuid: raise exceptions.CorruptStorage( f'The database has a repository UUID configured to {database_repository_uuid} ' f"but the disk-objectstore's is {repository_uuid}." )
[docs] def get_container(self) -> 'Container': """Return the disk-object store container. :returns: The disk-object store container configured for the repository path of the current profile. """ from disk_objectstore import Container from .backend import get_filepath_container return Container(get_filepath_container(self.profile))
[docs] def get_repository_uuid(self) -> str: """Return the UUID of the repository. :returns: The repository UUID. :raises: :class:`~aiida.common.exceptions.UnreachableStorage` if the UUID cannot be retrieved, which probably means that the repository is not initialised. """ try: return self.get_container().container_id except Exception as exception: raise exceptions.UnreachableStorage( f'Could not access disk-objectstore {self.get_container()}: {exception}' ) from exception
[docs] def initialise(self, reset: bool = False) -> bool: """Initialise the storage backend. This is typically used once when a new storage backed is created. If this method returns without exceptions the storage backend is ready for use. If the backend already seems initialised, this method is a no-op. :param reset: If ``true``, destroy the backend if it already exists including all of its data before recreating and initialising it. This is useful for example for test profiles that need to be reset before or after tests having run. :returns: ``True`` if the storage was initialised by the function call, ``False`` if it was already initialised. """ if reset: self.reset_repository() self.reset_database() initialised: bool = False if not self.is_initialised: self.initialise_repository() self.initialise_database() initialised = True # Call migrate in the case the storage was already initialised but not yet at the latest schema version. If it # was, then the following is a no-op anyway. self.migrate() return initialised
@property def is_initialised(self) -> bool: """Return whether the storage is initialised. This is the case if both the database and the repository are initialised. :returns: ``True`` if the storage is initialised, ``False`` otherwise. """ return self.is_repository_initialised and self.is_database_initialised @property def is_repository_initialised(self) -> bool: """Return whether the repository is initialised. :returns: ``True`` if the repository is initialised, ``False`` otherwise. """ return self.get_container().is_initialised @property def is_database_initialised(self) -> bool: """Return whether the database is initialised. This is the case if it contains the table that holds the schema version for alembic or Django. :returns: ``True`` if the database is initialised, ``False`` otherwise. """ return inspect(self.connection).has_table(self.alembic_version_tbl_name) or inspect(self.connection).has_table( )
[docs] def reset_repository(self) -> None: """Reset the repository by deleting all of its contents. This will also destroy the configuration and so in order to use it again, it will have to be reinitialised. """ import shutil try: shutil.rmtree(self.get_container().get_folder()) except FileNotFoundError: pass
[docs] def reset_database(self) -> None: """Reset the database by deleting all content from all tables. This will also destroy the settings table and so in order to use it again, it will have to be reinitialised. """ self.delete_all_tables(exclude_tables=[self.alembic_version_tbl_name])
[docs] def initialise_repository(self) -> None: """Initialise the repository.""" from import CONTAINER_DEFAULTS container = self.get_container() container.init_container(clear=True, **CONTAINER_DEFAULTS)
[docs] def initialise_database(self) -> None: """Initialise the database. This assumes that the database has no schema whatsoever and so the initial schema is created directly from the models at the current head version without migrating through all of them one by one. """ from import get_orm_metadata # setup the database # see:'initialising empty storage schema') get_orm_metadata().create_all(self._engine) repository_uuid = self.get_repository_uuid() # Create a "sync" between the database and repository, by saving its UUID in the settings table # this allows us to validate inconsistencies between the two self.connection.execute( insert(DbSetting).values(key=REPOSITORY_UUID_KEY, val=repository_uuid, description='Repository UUID') ) # finally, generate the version table, "stamping" it with the most recent revision with self._migration_context() as context: context.stamp(context.script, 'main@head') self.connection.commit()
[docs] def delete_all_tables(self, *, exclude_tables: list[str] | None = None) -> None: """Delete all tables of the current database schema. The tables are determined dynamically through reflection of the current schema version. Any other tables in the database that are not part of the schema should remain unaffected. :param exclude_tables: Optional list of table names that should not be deleted. """ exclude_tables = exclude_tables or [] if inspect(self.connection).has_table(self.alembic_version_tbl_name): metadata = MetaData() metadata.reflect(bind=self.connection) # The ``sorted_tables`` property returns the tables sorted by their foreign-key dependencies, with those # that are dependent on others first. Iterate over the list in reverse to ensure that the tables with # the independent rows are deleted first. for schema_table in reversed(metadata.sorted_tables): if in exclude_tables: continue self.connection.execute(schema_table.delete()) self.connection.commit()
[docs] def migrate(self) -> None: """Migrate the storage for this profile to the head version. :raises: :class:`~aiida.common.exceptions.UnreachableStorage` if the storage cannot be accessed. :raises: :class:`~aiida.common.exceptions.StorageMigrationError` if the storage is not initialised. """ # The database can be in one of a few states: # 1. Legacy django database -> we transfer the version to alembic, migrate to the head of the django branch, # reset the revision as one on the main branch, and then migrate to the head of the main branch # 2. Legacy sqlalchemy database -> we migrate to the head of the sqlalchemy branch, # reset the revision as one on the main branch, and then migrate to the head of the main branch # 3. Already on the main branch -> we migrate to the head of the main branch if not inspect(self.connection).has_table(self.alembic_version_tbl_name): if not inspect(self.connection).has_table( raise exceptions.StorageMigrationError('storage is uninitialised, cannot migrate.') # the database is a legacy django one, # so we need to copy the version from the 'django_migrations' table to the 'alembic_version' one legacy_version = self.get_schema_version_profile(check_legacy=True) if legacy_version is None: raise exceptions.StorageMigrationError( 'No schema version could be read from the database. ' "Check that either the 'alembic_version' or 'django_migrations' tables " 'are present and accessible, using e.g. `verdi devel run-sql "SELECT * FROM alembic_version"`' ) # the version should be of the format '00XX_description' version = f'django_{legacy_version[:4]}' with self._migration_context() as context: context.stamp(context.script, version) self.connection.commit() # now we can continue with the migration as normal else: version = self.get_schema_version_profile() # find what branch the current version is on branches = self._alembic_script().revision_map.get_revision(version).branch_labels if 'django' in branches or 'sqlalchemy' in branches: # migrate up to the top of the respective legacy branches if 'django' in branches:'Migrating to the head of the legacy django branch') self.migrate_up('django@head') elif 'sqlalchemy' in branches:'Migrating to the head of the legacy sqlalchemy branch') self.migrate_up('sqlalchemy@head') # now re-stamp with the comparable revision on the main branch with self._migration_context() as context: context._ensure_version_table(purge=True) context.stamp(context.script, 'main_0001') self.connection.commit() # finally migrate to the main head revision'Migrating to the head of the main branch') self.migrate_up('main@head') self.connection.commit()
[docs] def migrate_up(self, version: str) -> None: """Migrate the database up to a specific version. :param version: string with schema version to migrate to """ with self._alembic_connect() as config: upgrade(config, version)
[docs] def migrate_down(self, version: str) -> None: """Migrate the database down to a specific version. :param version: string with schema version to migrate to """ with self._alembic_connect() as config: downgrade(config, version)
[docs] @staticmethod def _alembic_config(): """Return an instance of an Alembic `Config`.""" dirpath = pathlib.Path(__file__).resolve().parent config = Config() config.set_main_option('script_location', str(dirpath / ALEMBIC_REL_PATH)) return config
[docs] @classmethod def _alembic_script(cls): """Return an instance of an Alembic `ScriptDirectory`.""" return ScriptDirectory.from_config(cls._alembic_config())
[docs] @contextlib.contextmanager def _alembic_connect(self) -> Iterator[Config]: """Context manager to return an instance of an Alembic configuration. The profiles's database connection is added in the `attributes` property, through which it can then also be retrieved, also in the `` file, which is run when the database is migrated. """ config = self._alembic_config() config.attributes['connection'] = self.connection config.attributes['aiida_profile'] = self.profile def _callback(step: MigrationInfo, **kwargs): """Callback to be called after a migration step is executed.""" from_rev = step.down_revision_ids[0] if step.down_revision_ids else '<base>''- {from_rev} -> {step.up_revision_id}') config.attributes['on_version_apply'] = _callback yield config
[docs] @contextlib.contextmanager def _migration_context(self) -> Iterator[MigrationContext]: """Context manager to return an instance of an Alembic migration context. This migration context will have been configured with the current database connection, which allows this context to be used to inspect the contents of the database, such as the current revision. """ with self._alembic_connect() as config: script = ScriptDirectory.from_config(config) with EnvironmentContext(config, script) as context: context.configure(context.config.attributes['connection']) yield context.get_context()
# the following are used for migration tests
[docs] @contextlib.contextmanager def session(self) -> Iterator[Session]: """Context manager to return a session for the database.""" session = Session(self._engine, future=True) try: yield session except Exception: session.rollback() raise finally: session.close()
[docs] def get_current_table(self, table_name: str) -> Any: """Return a table 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. """ base = automap_base() base.prepare(autoload_with=self.connection.engine) return getattr(base.classes, table_name)