From 351021164d00aa3a2a78b5b6e43e8a87a8553151 Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Mon, 15 May 2023 09:24:20 +0200 Subject: [PATCH] CLI: Add the command `verdi profile setup` This command uses the `DynamicEntryPointCommandGroup` to allow creating a new profile with any of the plugins registered in the `aiida.storage` group. Each storage plugin will typically require a different set of configuration parameters to initialize and connect to the storage. These are generated dynamically from the specification returned by the method `get_cli_options` defined on the `StorageBackend` base class. Each plugin implements the abstract `_get_cli_options` method which is called by the former and defines the configuration parameters of the plugin. The values passed to the plugin specific options are used to instantiate an instance of the storage class, registered under the chosen entry point which is then initialised. If successful, the new profile is stored in the `Config` and a default user is created and stored. After that, the profile is ready for use. --- aiida/cmdline/commands/cmd_profile.py | 39 ++++++++++++++- aiida/manage/configuration/__init__.py | 37 ++++++++++++++ aiida/orm/implementation/storage_backend.py | 13 +++++ aiida/storage/psql_dos/backend.py | 50 +++++++++++++++++++ aiida/storage/sqlite_temp/backend.py | 45 +++++++++++------ aiida/storage/sqlite_zip/backend.py | 19 ++++++- docs/source/reference/command_line.rst | 1 + tests/cmdline/commands/test_profile.py | 9 ++++ .../configuration/test_configuration.py | 13 ++++- tests/storage/sqlite/test_archive.py | 4 +- 10 files changed, 209 insertions(+), 21 deletions(-) diff --git a/aiida/cmdline/commands/cmd_profile.py b/aiida/cmdline/commands/cmd_profile.py index df7fc2331b..db2d017ab0 100644 --- a/aiida/cmdline/commands/cmd_profile.py +++ b/aiida/cmdline/commands/cmd_profile.py @@ -11,10 +11,12 @@ import click from aiida.cmdline.commands.cmd_verdi import verdi +from aiida.cmdline.groups import DynamicEntryPointCommandGroup from aiida.cmdline.params import arguments, options +from aiida.cmdline.params.options.commands import setup from aiida.cmdline.utils import defaults, echo from aiida.common import exceptions -from aiida.manage.configuration import get_config +from aiida.manage.configuration import Profile, create_profile, get_config @verdi.group('profile') @@ -22,6 +24,41 @@ def verdi_profile(): """Inspect and manage the configured profiles.""" +def command_create_profile(ctx: click.Context, storage_cls, non_interactive: bool, profile: Profile, **kwargs): # pylint: disable=unused-argument + """Create a new profile, initialise its storage and create a default user. + + :param ctx: The context of the CLI command. + :param storage_cls: The storage class obtained through loading the entry point from ``aiida.storage`` group. + :param non_interactive: Whether the command was invoked interactively or not. + :param profile: The profile instance. This is an empty ``Profile`` instance created by the command line argument + which currently only contains the selected profile name for the profile that is to be created. + :param kwargs: Arguments to initialise instance of the selected storage implementation. + """ + try: + profile = create_profile(ctx.obj.config, storage_cls, name=profile.name, **kwargs) + except (ValueError, TypeError, exceptions.EntryPointError, exceptions.StorageMigrationError) as exception: + echo.echo_critical(str(exception)) + + echo.echo_success(f'Created new profile `{profile.name}`.') + + +@verdi_profile.group( + 'setup', + cls=DynamicEntryPointCommandGroup, + command=command_create_profile, + entry_point_group='aiida.storage', + shared_options=[ + setup.SETUP_PROFILE(), + setup.SETUP_USER_EMAIL(), + setup.SETUP_USER_FIRST_NAME(), + setup.SETUP_USER_LAST_NAME(), + setup.SETUP_USER_INSTITUTION(), + ] +) +def profile_setup(): + """Set up a new profile.""" + + @verdi_profile.command('list') def profile_list(): """Display a list of all available profiles.""" diff --git a/aiida/manage/configuration/__init__.py b/aiida/manage/configuration/__init__.py index fe92492bc1..22d3c8d31b 100644 --- a/aiida/manage/configuration/__init__.py +++ b/aiida/manage/configuration/__init__.py @@ -184,6 +184,43 @@ def profile_context(profile: Optional[str] = None, allow_switch=False) -> 'Profi manager.load_profile(current_profile, allow_switch=True) +def create_profile( + config: Config, + storage_cls, + *, + name: str, + email: str, + first_name: Optional[str] = None, + last_name: Optional[str] = None, + institution: Optional[str] = None, + **kwargs +) -> Profile: + """Create a new profile, initialise its storage and create a default user. + + :param config: The config instance. + :param storage_cls: The storage class obtained through loading the entry point from ``aiida.storage`` group. + :param name: Name of the profile. + :param email: Email for the default user. + :param first_name: First name for the default user. + :param last_name: Last name for the default user. + :param institution: Institution for the default user. + :param kwargs: Arguments to initialise instance of the selected storage implementation. + """ + from aiida.orm import User + + storage_config = {key: kwargs[key] for key in storage_cls.get_cli_options().keys() if key in kwargs} + profile: Profile = config.create_profile(name=name, storage_cls=storage_cls, storage_config=storage_config) + + with profile_context(profile.name, allow_switch=True): + user = User(email=email, first_name=first_name, last_name=last_name, institution=institution).store() + profile.default_user_email = user.email + + config.update_profile(profile) + config.store() + + return profile + + def reset_config(): """Reset the globally loaded config. diff --git a/aiida/orm/implementation/storage_backend.py b/aiida/orm/implementation/storage_backend.py index bc99001ffb..b1a19d73bb 100644 --- a/aiida/orm/implementation/storage_backend.py +++ b/aiida/orm/implementation/storage_backend.py @@ -8,7 +8,10 @@ # For further information please visit http://www.aiida.net # ########################################################################### """Generic backend related objects""" +from __future__ import annotations + import abc +import collections from typing import TYPE_CHECKING, Any, ContextManager, List, Optional, Sequence, TypeVar, Union if TYPE_CHECKING: @@ -89,6 +92,16 @@ def migrate(cls, profile: 'Profile') -> None: :raises: :class:`~aiida.common.exceptions.StorageMigrationError` if the storage is not initialised. """ + @classmethod + def get_cli_options(cls) -> collections.OrderedDict: + """Return the CLI options that would allow to create an instance of this class.""" + return collections.OrderedDict(cls._get_cli_options()) + + @classmethod + @abc.abstractmethod + def _get_cli_options(cls) -> dict[str, Any]: + """Return the CLI options that would allow to create an instance of this class.""" + @abc.abstractmethod def __init__(self, profile: 'Profile') -> None: """Initialize the backend, for this profile. diff --git a/aiida/storage/psql_dos/backend.py b/aiida/storage/psql_dos/backend.py index c425fcfe7e..0a5a9ee560 100644 --- a/aiida/storage/psql_dos/backend.py +++ b/aiida/storage/psql_dos/backend.py @@ -102,6 +102,56 @@ def migrator_context(cls, profile: Profile): finally: migrator.close() + @classmethod + def create_config(cls, **kwargs): + """Create a configuration dictionary based on the CLI options that can be used to initialize an instance.""" + return {key: kwargs[key] for key in cls._get_cli_options()} + + @classmethod + def _get_cli_options(cls) -> dict: + """Return the CLI options that would allow to create an instance of this class.""" + return { + 'database_engine': { + 'required': True, + 'type': str, + 'prompt': 'Postgresql engine', + 'default': 'postgresql_psycopg2', + 'help': 'The engine to use to connect to the database.', + }, + 'database_hostname': { + 'required': True, + 'type': str, + 'prompt': 'Postgresql hostname', + 'default': 'localhost', + 'help': 'The hostname of the PostgreSQL server.', + }, + 'database_port': { + 'required': True, + 'type': int, + 'prompt': 'Postgresql port', + 'default': '5432', + 'help': 'The port of the PostgreSQL server.', + }, + 'database_username': { + 'required': True, + 'type': str, + 'prompt': 'Postgresql username', + 'help': 'The username with which to connect to the PostgreSQL server.', + }, + 'database_password': { + 'required': True, + 'type': str, + 'prompt': 'Postgresql password', + 'help': 'The password with which to connect to the PostgreSQL server.', + }, + 'database_name': { + 'required': True, + 'type': str, + 'prompt': 'Postgresql database name', + 'help': 'The name of the database in the PostgreSQL server.', + } + } + def __init__(self, profile: Profile) -> None: super().__init__(profile) diff --git a/aiida/storage/sqlite_temp/backend.py b/aiida/storage/sqlite_temp/backend.py index c25a9d7777..81b3a4a26e 100644 --- a/aiida/storage/sqlite_temp/backend.py +++ b/aiida/storage/sqlite_temp/backend.py @@ -12,18 +12,18 @@ from contextlib import contextmanager, nullcontext import functools -from functools import cached_property import hashlib import os from pathlib import Path import shutil +from tempfile import mkdtemp from typing import Any, BinaryIO, Iterator, Sequence from sqlalchemy import column, insert, update from sqlalchemy.orm import Session from aiida.common.exceptions import ClosedStorage, IntegrityError -from aiida.manage import Profile +from aiida.manage.configuration import Profile from aiida.orm.entities import EntityTypes from aiida.orm.implementation import BackendEntity, StorageBackend from aiida.repository.backend.sandbox import SandboxRepositoryBackend @@ -49,7 +49,7 @@ def create_profile( default_user_email='user@email.com', options: dict | None = None, debug: bool = False, - repo_path: str | Path | None = None, + filepath: str | Path | None = None, ) -> Profile: """Create a new profile instance for this backend, from the path to the zip file.""" return Profile( @@ -58,8 +58,8 @@ def create_profile( 'storage': { 'backend': 'core.sqlite_temp', 'config': { + 'filepath': filepath, 'debug': debug, - 'repo_path': repo_path, } }, 'process_control': { @@ -70,6 +70,23 @@ def create_profile( } ) + @classmethod + def create_config(cls, filepath: str | None = None): + """Create a configuration dictionary based on the CLI options that can be used to initialize an instance.""" + return {'filepath': filepath or mkdtemp()} + + @classmethod + def _get_cli_options(cls) -> dict: + """Return the CLI options that would allow to create an instance of this class.""" + return { + 'filepath': { + 'required': False, + 'type': str, + 'prompt': 'Temporary directory', + 'help': 'Temporary directory in which to store data for this backend.', + } + } + @classmethod def version_head(cls) -> str: return get_schema_version_head() @@ -89,7 +106,7 @@ def migrate(cls, profile: Profile): def __init__(self, profile: Profile): super().__init__(profile) self._session: Session | None = None - self._repo: SandboxShaRepositoryBackend | None = None + self._repo: SandboxShaRepositoryBackend = SandboxShaRepositoryBackend(profile.storage_config['filepath']) self._globals: dict[str, tuple[Any, str | None]] = {} self._closed = False self.get_session() # load the database on initialization @@ -135,10 +152,6 @@ def get_session(self) -> Session: def get_repository(self) -> SandboxShaRepositoryBackend: if self._closed: raise ClosedStorage(str(self)) - if self._repo is None: - # to-do this does not seem to be removing the folder on garbage collection? - repo_path = self.profile.storage_config.get('repo_path') - self._repo = SandboxShaRepositoryBackend(filepath=Path(repo_path) if repo_path else None) return self._repo @property @@ -175,31 +188,31 @@ def get_backend_entity(self, model) -> BackendEntity: """Return the backend entity that corresponds to the given Model instance.""" return orm.get_backend_entity(model, self) - @cached_property + @functools.cached_property def authinfos(self): return orm.SqliteAuthInfoCollection(self) - @cached_property + @functools.cached_property def comments(self): return orm.SqliteCommentCollection(self) - @cached_property + @functools.cached_property def computers(self): return orm.SqliteComputerCollection(self) - @cached_property + @functools.cached_property def groups(self): return orm.SqliteGroupCollection(self) - @cached_property + @functools.cached_property def logs(self): return orm.SqliteLogCollection(self) - @cached_property + @functools.cached_property def nodes(self): return orm.SqliteNodeCollection(self) - @cached_property + @functools.cached_property def users(self): return orm.SqliteUserCollection(self) diff --git a/aiida/storage/sqlite_zip/backend.py b/aiida/storage/sqlite_zip/backend.py index dd58451bf1..7043ff0fd5 100644 --- a/aiida/storage/sqlite_zip/backend.py +++ b/aiida/storage/sqlite_zip/backend.py @@ -89,6 +89,23 @@ def create_profile(path: str | Path, options: dict | None = None) -> Profile: } ) + @classmethod + def create_config(cls, filepath: str): + """Create a configuration dictionary based on the CLI options that can be used to initialize an instance.""" + return {'path': filepath} + + @classmethod + def _get_cli_options(cls) -> dict: + """Return the CLI options that would allow to create an instance of this class.""" + return { + 'filepath': { + 'required': True, + 'type': str, + 'prompt': 'Filepath of the archive', + 'help': 'Filepath of the archive in which to store data for this backend.', + } + } + @classmethod def version_profile(cls, profile: Profile) -> Optional[str]: return read_version(profile.storage_config['path'], search_limit=None) @@ -427,5 +444,5 @@ def list_objects(self) -> Iterable[str]: def open(self, key: str) -> Iterator[BinaryIO]: if not self._path.joinpath(key).is_file(): raise FileNotFoundError(f'object with key `{key}` does not exist.') - with self._path.joinpath(key).open('rb') as handle: + with self._path.joinpath(key).open('rb', encoding='utf-8') as handle: yield handle diff --git a/docs/source/reference/command_line.rst b/docs/source/reference/command_line.rst index c3ae6e2f7c..df68843a24 100644 --- a/docs/source/reference/command_line.rst +++ b/docs/source/reference/command_line.rst @@ -356,6 +356,7 @@ Below is a list with all available subcommands. delete Delete one or more profiles. list Display a list of all available profiles. setdefault Set a profile as the default one. + setup Set up a new profile. show Show details for a profile. diff --git a/tests/cmdline/commands/test_profile.py b/tests/cmdline/commands/test_profile.py index 7b2d6f20b5..2cd8201d6e 100644 --- a/tests/cmdline/commands/test_profile.py +++ b/tests/cmdline/commands/test_profile.py @@ -131,3 +131,12 @@ def test_delete(run_cli_command, mock_profiles, pg_test_cluster): result = run_cli_command(cmd_profile.profile_list, use_subprocess=False) assert profile_list[2] not in result.output assert profile_list[3] not in result.output + + +def test_setup(run_cli_command, isolated_config, tmp_path): + """Test the ``verdi profile setup`` command.""" + profile_name = 'temp-profile' + options = ['core.sqlite_temp', '-n', '--filepath', str(tmp_path), '--profile', profile_name] + result = run_cli_command(cmd_profile.profile_setup, options, use_subprocess=False) + assert f'Created new profile `{profile_name}`.' in result.output + assert profile_name in isolated_config.profile_names diff --git a/tests/manage/configuration/test_configuration.py b/tests/manage/configuration/test_configuration.py index de9fd3a417..e363bd9a60 100644 --- a/tests/manage/configuration/test_configuration.py +++ b/tests/manage/configuration/test_configuration.py @@ -3,8 +3,19 @@ import pytest import aiida -from aiida.manage.configuration import get_profile, profile_context +from aiida.manage.configuration import Profile, create_profile, get_profile, profile_context from aiida.manage.manager import get_manager +from aiida.storage.sqlite_temp.backend import SqliteTempBackend + + +def test_create_profile(isolated_config, tmp_path): + """Test :func:`aiida.manage.configuration.tools.create_profile`.""" + profile_name = 'testing' + profile = create_profile( + isolated_config, SqliteTempBackend, name=profile_name, email='test@localhost', filepath=str(tmp_path) + ) + assert isinstance(profile, Profile) + assert profile_name in isolated_config.profile_names def test_check_version_release(monkeypatch, capsys, isolated_config): diff --git a/tests/storage/sqlite/test_archive.py b/tests/storage/sqlite/test_archive.py index b370b67966..6169876c08 100644 --- a/tests/storage/sqlite/test_archive.py +++ b/tests/storage/sqlite/test_archive.py @@ -12,7 +12,7 @@ def test_basic(tmp_path): filename = Path(tmp_path / 'export.aiida') # generate a temporary backend - profile1 = SqliteTempBackend.create_profile(repo_path=str(tmp_path / 'repo1')) + profile1 = SqliteTempBackend.create_profile(filepath=str(tmp_path / 'repo1')) backend1 = SqliteTempBackend(profile1) # add simple node @@ -30,7 +30,7 @@ def test_basic(tmp_path): create_archive(None, backend=backend1, filename=filename) # create a new temporary backend and import - profile2 = SqliteTempBackend.create_profile(repo_path=str(tmp_path / 'repo2')) + profile2 = SqliteTempBackend.create_profile(filepath=str(tmp_path / 'repo2')) backend2 = SqliteTempBackend(profile2) import_archive(filename, backend=backend2)