aiida.storage.psql_dos package

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

Subpackages

• aiida.storage.psql_dos.migrations package
  • Subpackages
    • aiida.storage.psql_dos.migrations.utils package
      • Submodules
  • Submodules
• aiida.storage.psql_dos.models package
  • Submodules
• aiida.storage.psql_dos.orm package
  • Subpackages
    • aiida.storage.psql_dos.orm.querybuilder package
      • Submodules
  • Submodules

Submodules

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

class aiida.storage.psql_dos.alembic_cli.AlembicRunner

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__': {'profile': 'Profile | None'}})

__init__() → None

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

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)

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

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

Return a string showing connection details for this instance.

_abc_impl = <_abc._abc_data object>

_clear() → None

Clear the storage, removing all data.

Warning

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

_default_user: Optional['User']

static _get_mapper_from_entity(entity_type: aiida.orm.entities.EntityTypes, with_pk: bool)

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

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]

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

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

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

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

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]

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

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

Return the object repository configured for this backend.

get_session() → sqlalchemy.orm.session.Session

Return an SQLAlchemy session bound to the current thread.

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

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.

classmethod initialise(profile: aiida.manage.configuration.profile.Profile, 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.

Parameters

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.

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

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

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 uninitialised, this method will raise an exception.

Raises

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

Raises

StorageMigrationError if the storage is not initialised.

migrator

alias of aiida.storage.psql_dos.migrator.PsqlDosMigrator

classmethod migrator_context(profile: aiida.manage.configuration.profile.Profile)

property nodes

Return the collection of nodes

query()

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

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]

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

Return the head schema version of this storage backend type.

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

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

aiida.storage.psql_dos.backend.get_filepath_container(profile: aiida.manage.configuration.profile.Profile) → pathlib.Path

Return the filepath of the disk-object store container.

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.PsqlDosMigrator(profile: aiida.manage.configuration.profile.Profile)

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 0x7ff85ab6d120; django_migrations>, '__init__': <function PsqlDosMigrator.__init__>, 'close': <function PsqlDosMigrator.close>, 'connection': <property object>, 'get_schema_versions': <classmethod(<function PsqlDosMigrator.get_schema_versions>)>, 'get_schema_version_head': <classmethod(<function PsqlDosMigrator.get_schema_version_head>)>, 'get_schema_version_profile': <function PsqlDosMigrator.get_schema_version_profile>, 'validate_storage': <function PsqlDosMigrator.validate_storage>, 'get_container': <function PsqlDosMigrator.get_container>, 'get_repository_uuid': <function PsqlDosMigrator.get_repository_uuid>, 'initialise': <function PsqlDosMigrator.initialise>, 'is_initialised': <property object>, 'is_repository_initialised': <property object>, 'is_database_initialised': <property object>, 'reset_repository': <function PsqlDosMigrator.reset_repository>, 'reset_database': <function PsqlDosMigrator.reset_database>, 'initialise_repository': <function PsqlDosMigrator.initialise_repository>, 'initialise_database': <function PsqlDosMigrator.initialise_database>, 'delete_all_tables': <function PsqlDosMigrator.delete_all_tables>, 'migrate': <function PsqlDosMigrator.migrate>, 'migrate_up': <function PsqlDosMigrator.migrate_up>, 'migrate_down': <function PsqlDosMigrator.migrate_down>, '_alembic_config': <staticmethod(<function PsqlDosMigrator._alembic_config>)>, '_alembic_script': <classmethod(<function PsqlDosMigrator._alembic_script>)>, '_alembic_connect': <function PsqlDosMigrator._alembic_connect>, '_migration_context': <function PsqlDosMigrator._migration_context>, 'session': <function PsqlDosMigrator.session>, 'get_current_table': <function PsqlDosMigrator.get_current_table>, '__dict__': <attribute '__dict__' of 'PsqlDosMigrator' objects>, '__weakref__': <attribute '__weakref__' of 'PsqlDosMigrator' objects>, '__annotations__': {}})

__init__(profile: aiida.manage.configuration.profile.Profile) → None

__module__ = 'aiida.storage.psql_dos.migrator'

__weakref__

list of weak references to the object (if defined)

static _alembic_config()

Return an instance of an Alembic Config.

_alembic_connect() → Iterator[alembic.config.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 env.py file, which is run when the database is migrated.

classmethod _alembic_script()

Return an instance of an Alembic ScriptDirectory.

_migration_context() → Iterator[alembic.runtime.migration.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.

alembic_version_tbl_name = 'alembic_version'

close() → None

Close the connection if it was opened and dispose of the engine.

property connection

Return the connection to the database.

Will automatically create the engine and open an connection if not already opened in a previous call.

Returns

Open connection to the database.

Raises

aiida.common.exceptions.UnreachableStorage if connecting to the database fails.

delete_all_tables(*, exclude_tables: Optional[list[str]] = 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.

Parameters

exclude_tables – Optional list of table names that should not be deleted.

django_version_table = <sqlalchemy.sql.selectable.TableClause at 0x7ff85ab6d120; django_migrations>

get_container() → Container

Return the disk-object store container.

Returns

The disk-object store container configured for the repository path of the current profile.

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

get_repository_uuid() → str

Return the UUID of the repository.

Returns

The repository UUID.

Raises

UnreachableStorage if the UUID cannot be retrieved, which probably means that the repository is not initialised.

classmethod get_schema_version_head() → str

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

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

classmethod get_schema_versions() → Dict[str, str]

Return all available schema versions (oldest to latest).

Returns

schema version -> description

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

Parameters

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.

initialise_database() → 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.

initialise_repository() → None

Initialise the repository.

property is_database_initialised: 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.

property is_initialised: 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.

property is_repository_initialised: bool

Return whether the repository is initialised.

Returns

True if the repository is initialised, False otherwise.

migrate() → None

Migrate the storage for this profile to the head version.

Raises

UnreachableStorage if the storage cannot be accessed.

Raises

StorageMigrationError if the storage is not initialised.

migrate_down(version: str) → None

Migrate the database down to a specific version.

Parameters

version – string with schema version to migrate to

migrate_up(version: str) → None

Migrate the database up to a specific version.

Parameters

version – string with schema version to migrate to

reset_database() → 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.

reset_repository() → 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.

session() → Iterator[sqlalchemy.orm.session.Session]

Context manager to return a session for the database.

validate_storage() → 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

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

Bases: TypedDict

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.', '__orig_bases__': (<function TypedDict>,), '__dict__': <attribute '__dict__' of 'PsqlConfig' objects>, '__weakref__': <attribute '__weakref__' of 'PsqlConfig' objects>, '__required_keys__': frozenset(), '__optional_keys__': frozenset({'database_password', 'database_name', 'engine_kwargs', 'database_username', 'database_port', 'database_hostname'}), '__total__': False})

__module__ = 'aiida.storage.psql_dos.utils'

__optional_keys__ = frozenset({'database_hostname', 'database_name', 'database_password', 'database_port', 'database_username', 'engine_kwargs'})

__orig_bases__ = (<function TypedDict>,)

__required_keys__ = frozenset({})

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

Create scoped SQLAlchemy session factory

aiida.storage.psql_dos.utils.create_sqlalchemy_engine(config: aiida.storage.psql_dos.utils.PsqlConfig)

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)

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)

Return the transitive closure table template

aiida.storage.psql_dos.utils.install_tc(session)

Install the transitive closure table with SqlAlchemy.