aiida.storage.psql_dos package#

Module with implementation of the storage backend using PostGreSQL and the disk-objectstore.

Submodules#

Simple wrapper around the alembic command line tool that first loads an AiiDA profile.

class aiida.storage.psql_dos.alembic_cli.AlembicRunner[source]#

Bases: object

Wrapper around the alembic command line tool that first loads an AiiDA profile.

__dict__ = mappingproxy({'__module__': 'aiida.storage.psql_dos.alembic_cli', '__doc__': 'Wrapper around the alembic command line tool that first loads an AiiDA profile.', '__init__': <function AlembicRunner.__init__>, 'execute_alembic_command': <function AlembicRunner.execute_alembic_command>, '__dict__': <attribute '__dict__' of 'AlembicRunner' objects>, '__weakref__': <attribute '__weakref__' of 'AlembicRunner' objects>, '__annotations__': {}})#
__init__() None[source]#
__module__ = 'aiida.storage.psql_dos.alembic_cli'#
__weakref__#

list of weak references to the object (if defined)

execute_alembic_command(command_name, connect=True, **kwargs)[source]#

Execute an Alembic CLI command.

Parameters
  • command_name – the sub command name

  • kwargs – parameters to pass to the command

SqlAlchemy implementation of aiida.orm.implementation.backends.Backend.

class aiida.storage.psql_dos.backend.PsqlDosBackend(profile: aiida.manage.configuration.profile.Profile)[source]#

Bases: aiida.orm.implementation.storage_backend.StorageBackend

An AiiDA storage backend that stores data in a PostgreSQL database and disk-objectstore repository.

Note, there were originally two such backends, sqlalchemy and django. The django backend was removed, to consolidate access to this storage.

__abstractmethods__ = frozenset({})#
__init__(profile: aiida.manage.configuration.profile.Profile) None[source]#

Initialize the backend, for this profile.

Raises

~aiida.common.exceptions.UnreachableStorage if the storage cannot be accessed

Raises

~aiida.common.exceptions.IncompatibleStorageSchema if the profile’s storage schema is not at the latest version (and thus should be migrated)

Raises
raises

aiida.common.exceptions.CorruptStorage if the storage is internally inconsistent

__module__ = 'aiida.storage.psql_dos.backend'#
__str__() str[source]#

Return a string showing connection details for this instance.

_abc_impl = <_abc_data object>#
_clear(recreate_user: bool = True) None[source]#

Clear the storage, removing all data.

Warning

This is a destructive operation, and should only be used for testing purposes.

Parameters

recreate_user – Re-create the default User for the profile, after clearing the storage.

_default_user: Optional[User]#
static _get_mapper_from_entity(entity_type: aiida.orm.entities.EntityTypes, with_pk: bool)[source]#

Return the Sqlalchemy mapper and fields corresponding to the given entity.

Parameters

with_pk – if True, the fields returned will include the primary key

_initialise_session()[source]#

Initialise the SQLAlchemy session factory.

Only one session factory is ever associated with a given class instance, i.e. once the instance is closed, it cannot be reopened.

The session factory, returns a session that is bound to the current thread. Multi-thread support is currently required by the REST API. Although, in the future, we may want to move the multi-thread handling to higher in the AiiDA stack.

property authinfos#

Return the collection of authorisation information objects

bulk_insert(entity_type: aiida.orm.entities.EntityTypes, rows: List[dict], allow_defaults: bool = False) List[int][source]#

Insert a list of entities into the database, directly into a backend transaction.

Parameters
  • entity_type – The type of the entity

  • data – A list of dictionaries, containing all fields of the backend model, except the id field (a.k.a primary key), which will be generated dynamically

  • allow_defaults – If False, assert that each row contains all fields (except primary key(s)), otherwise, allow default values for missing fields.

Raises

IntegrityError if the keys in a row are not a subset of the columns in the table

Returns

The list of generated primary keys for the entities

bulk_update(entity_type: aiida.orm.entities.EntityTypes, rows: List[dict]) None[source]#

Update a list of entities in the database, directly with a backend transaction.

Parameters
  • entity_type – The type of the entity

  • data – A list of dictionaries, containing fields of the backend model to update, and the id field (a.k.a primary key)

Raises

IntegrityError if the keys in a row are not a subset of the columns in the table

close() None[source]#

Close the storage access.

property comments#

Return the collection of comments

property computers#

Return the collection of computers

delete_nodes_and_connections(pks_to_delete: Sequence[int]) None[source]#

Delete all nodes corresponding to pks in the input and any links to/from them.

This method is intended to be used within a transaction context.

Parameters

pks_to_delete – a sequence of node pks to delete

Raises

AssertionError if a transaction is not active

get_backend_entity(model: sqlalchemy.orm.decl_api.Model) aiida.orm.implementation.entities.BackendEntity[source]#

Return the backend entity that corresponds to the given Model instance

Parameters

model – the ORM model instance to promote to a backend instance

Returns

the backend entity corresponding to the given model

get_global_variable(key: str) Union[None, str, int, float][source]#

Return a global variable from the storage.

Parameters

key – the key of the setting

Raises

KeyError if the setting does not exist

get_info(detailed: bool = False) dict[source]#

Return general information on the storage.

Parameters

detailed – flag to request more detailed information about the content of the storage.

Returns

a nested dict with the relevant information.

get_repository() DiskObjectStoreRepositoryBackend[source]#

Return the object repository configured for this backend.

get_session() sqlalchemy.orm.session.Session[source]#

Return an SQLAlchemy session bound to the current thread.

get_unreferenced_keyset(check_consistency: bool = True) Set[str][source]#

Returns the keyset of objects that exist in the repository but are not tracked by AiiDA.

This should be all the soft-deleted files.

Parameters

check_consistency – toggle for a check that raises if there are references in the database with no actual object in the underlying repository.

Returns

a set with all the objects in the underlying repository that are not referenced in the database.

property groups#

Return the collection of groups

property in_transaction: bool#

Return whether a transaction is currently active.

property is_closed: bool#

Return whether the storage is closed.

property logs#

Return the collection of logs

maintain(full: bool = False, dry_run: bool = False, **kwargs) None[source]#

Perform maintenance tasks on the storage.

If full == True, then this method may attempt to block the profile associated with the storage to guarantee the safety of its procedures. This will not only prevent any other subsequent process from accessing that profile, but will also first check if there is already any process using it and raise if that is the case. The user will have to manually stop any processes that is currently accessing the profile themselves or wait for it to finish on its own.

Parameters
  • full – flag to perform operations that require to stop using the profile to be maintained.

  • dry_run – flag to only print the actions that would be taken without actually executing them.

classmethod migrate(profile: aiida.manage.configuration.profile.Profile) None[source]#

Migrate the storage of a profile to the latest schema version.

If the schema version is already the latest version, this method does nothing. If the storage is empty/uninitialised, then it will be initialised at head.

Raises

~aiida.common.exceptions.UnreachableStorage if the storage cannot be accessed

migrator#

alias of aiida.storage.psql_dos.migrator.PsqlDostoreMigrator

property nodes#

Return the collection of nodes

query()[source]#

Return an instance of a query builder implementation for this backend

set_global_variable(key: str, value: Union[None, str, int, float], description: Optional[str] = None, overwrite=True) None[source]#

Set a global variable in the storage.

Parameters
  • key – the key of the setting

  • value – the value of the setting

  • description – the description of the setting (optional)

  • overwrite – if True, overwrite the setting if it already exists

Raises

ValueError if the key already exists and overwrite is False

transaction() Iterator[sqlalchemy.orm.session.Session][source]#

Open a transaction to be used as a context manager.

If there is an exception within the context then the changes will be rolled back and the state will be as before entering. Transactions can be nested.

property users#

Return the collection of users

classmethod version_head() str[source]#

Return the head schema version of this storage backend type.

classmethod version_profile(profile: aiida.manage.configuration.profile.Profile) Optional[str][source]#

Return the schema version of the given profile’s storage, or None for empty/uninitialised storage.

Raises

~aiida.common.exceptions.UnreachableStorage if the storage cannot be accessed

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!

class aiida.storage.psql_dos.migrator.PsqlDostoreMigrator(profile: aiida.manage.configuration.profile.Profile)[source]#

Bases: object

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)

__dict__ = mappingproxy({'__module__': 'aiida.storage.psql_dos.migrator', '__doc__': 'Class for validating and migrating `psql_dos` storage instances.\n\n    .. important:: This class should only be accessed via the storage backend class (apart from for test purposes)\n    ', 'alembic_version_tbl_name': 'alembic_version', 'django_version_table': <sqlalchemy.sql.selectable.TableClause at 0x7fb89892bb80; django_migrations>, '__init__': <function PsqlDostoreMigrator.__init__>, 'get_schema_versions': <classmethod object>, 'get_schema_version_head': <classmethod object>, '_connection_context': <function PsqlDostoreMigrator._connection_context>, 'get_schema_version_profile': <function PsqlDostoreMigrator.get_schema_version_profile>, 'validate_storage': <function PsqlDostoreMigrator.validate_storage>, 'initialise': <function PsqlDostoreMigrator.initialise>, 'migrate': <function PsqlDostoreMigrator.migrate>, 'migrate_up': <function PsqlDostoreMigrator.migrate_up>, 'migrate_down': <function PsqlDostoreMigrator.migrate_down>, '_alembic_config': <staticmethod object>, '_alembic_script': <classmethod object>, '_alembic_connect': <function PsqlDostoreMigrator._alembic_connect>, '_migration_context': <function PsqlDostoreMigrator._migration_context>, 'session': <function PsqlDostoreMigrator.session>, 'get_current_table': <function PsqlDostoreMigrator.get_current_table>, '__dict__': <attribute '__dict__' of 'PsqlDostoreMigrator' objects>, '__weakref__': <attribute '__weakref__' of 'PsqlDostoreMigrator' objects>, '__annotations__': {}})#
__init__(profile: aiida.manage.configuration.profile.Profile) None[source]#
__module__ = 'aiida.storage.psql_dos.migrator'#
__weakref__#

list of weak references to the object (if defined)

static _alembic_config()[source]#

Return an instance of an Alembic Config.

_alembic_connect(_connection: Optional[sqlalchemy.future.engine.Connection] = None) Iterator[alembic.config.Config][source]#

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 env.py file, which is run when the database is migrated.

classmethod _alembic_script()[source]#

Return an instance of an Alembic ScriptDirectory.

_connection_context(connection: Optional[sqlalchemy.future.engine.Connection] = None) AbstractContextManager[sqlalchemy.future.engine.Connection][source]#

Return a context manager, with a connection to the database.

Raises

UnreachableStorage if the database connection fails

_migration_context(_connection: Optional[sqlalchemy.future.engine.Connection] = None) Iterator[alembic.runtime.migration.MigrationContext][source]#

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.

alembic_version_tbl_name = 'alembic_version'#
django_version_table = <sqlalchemy.sql.selectable.TableClause at 0x7fb89892bb80; django_migrations>#
get_current_table(table_name: str) sqlalchemy.sql.schema.Table[source]#

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.

classmethod get_schema_version_head() str[source]#

Return the head schema version for this storage, i.e. the latest schema this storage can be migrated to.

get_schema_version_profile(_connection: Optional[sqlalchemy.future.engine.Connection] = None, check_legacy=False) Optional[str][source]#

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.

classmethod get_schema_versions() Dict[str, str][source]#

Return all available schema versions (oldest to latest).

Returns

schema version -> description

initialise() None[source]#

Generate the initial storage schema for this profile, from the ORM models.

migrate() None[source]#

Migrate the storage for this profile to the head version.

Raises

UnreachableStorage if the storage cannot be accessed

migrate_down(version: str) None[source]#

Migrate the database down to a specific version.

Parameters

version – string with schema version to migrate to

migrate_up(version: str) None[source]#

Migrate the database up to a specific version.

Parameters

version – string with schema version to migrate to

session() Iterator[sqlalchemy.orm.session.Session][source]#

Context manager to return a session for the database.

validate_storage() None[source]#

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

aiida.common.exceptions.UnreachableStorage if the storage cannot be connected to

Raises

aiida.common.exceptions.IncompatibleStorageSchema if the storage is not compatible with the code API.

Raises

aiida.common.exceptions.CorruptStorage if the repository ID is not equal to the UUID set in thedatabase.

Utility functions specific to the SqlAlchemy backend.

class aiida.storage.psql_dos.utils.PsqlConfig(*args, **kwargs)[source]#

Bases: dict

Configuration to connect to a PostgreSQL database.

__annotations__ = {'database_hostname': <class 'str'>, 'database_name': <class 'str'>, 'database_password': <class 'str'>, 'database_port': <class 'int'>, 'database_username': <class 'str'>, 'engine_kwargs': <class 'dict'>}#
__dict__ = mappingproxy({'__module__': 'aiida.storage.psql_dos.utils', '__annotations__': {'database_hostname': <class 'str'>, 'database_port': <class 'int'>, 'database_username': <class 'str'>, 'database_password': <class 'str'>, 'database_name': <class 'str'>, 'engine_kwargs': <class 'dict'>}, '__doc__': 'Configuration to connect to a PostgreSQL database.', '__new__': <staticmethod object>, '__dict__': <attribute '__dict__' of 'PsqlConfig' objects>, '__weakref__': <attribute '__weakref__' of 'PsqlConfig' objects>, '__total__': False})#
__module__ = 'aiida.storage.psql_dos.utils'#
static __new__(cls, /, *args, **kwargs)#
__total__ = False#
__weakref__#

list of weak references to the object (if defined)

database_hostname: str#
database_name: str#
database_password: str#
database_port: int#
database_username: str#
engine_kwargs: dict#

keyword argument that will be passed on to the SQLAlchemy engine.

aiida.storage.psql_dos.utils.create_scoped_session_factory(engine, **kwargs)[source]#

Create scoped SQLAlchemy session factory

aiida.storage.psql_dos.utils.create_sqlalchemy_engine(config: aiida.storage.psql_dos.utils.PsqlConfig)[source]#

Create SQLAlchemy engine (to be used for QueryBuilder queries)

Parameters

kwargs – keyword arguments that will be passed on to sqlalchemy.create_engine. See https://docs.sqlalchemy.org/en/13/core/engines.html?highlight=create_engine#sqlalchemy.create_engine for more info.

aiida.storage.psql_dos.utils.flag_modified(instance, key)[source]#

Wrapper around sqlalchemy.orm.attributes.flag_modified to correctly dereference utils.ModelWrapper

Since SqlAlchemy 1.2.12 (and maybe earlier but not in 1.0.19) the flag_modified function will check that the key is actually present in the instance or it will except. If we pass a model instance, wrapped in the ModelWrapper the call will raise an InvalidRequestError. In this function that wraps the flag_modified of SqlAlchemy, we derefence the model instance if the passed instance is actually wrapped in the ModelWrapper.

aiida.storage.psql_dos.utils.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)[source]#

Return the transitive closure table template

aiida.storage.psql_dos.utils.install_tc(session)[source]#

Install the transitive closure table with SqlAlchemy.