Skip to content

Commit

Permalink
CLI: Add the command verdi profile setup
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
sphuber committed Nov 9, 2023
1 parent ae7abe8 commit 3510211
Show file tree
Hide file tree
Showing 10 changed files with 209 additions and 21 deletions.
39 changes: 38 additions & 1 deletion aiida/cmdline/commands/cmd_profile.py
Expand Up @@ -11,17 +11,54 @@
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')
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."""
Expand Down
37 changes: 37 additions & 0 deletions aiida/manage/configuration/__init__.py
Expand Up @@ -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.
Expand Down
13 changes: 13 additions & 0 deletions aiida/orm/implementation/storage_backend.py
Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down
50 changes: 50 additions & 0 deletions aiida/storage/psql_dos/backend.py
Expand Up @@ -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)

Expand Down
45 changes: 29 additions & 16 deletions aiida/storage/sqlite_temp/backend.py
Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -58,8 +58,8 @@ def create_profile(
'storage': {
'backend': 'core.sqlite_temp',
'config': {
'filepath': filepath,
'debug': debug,
'repo_path': repo_path,
}
},
'process_control': {
Expand All @@ -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()
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down
19 changes: 18 additions & 1 deletion aiida/storage/sqlite_zip/backend.py
Expand Up @@ -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)
Expand Down Expand Up @@ -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
1 change: 1 addition & 0 deletions docs/source/reference/command_line.rst
Expand Up @@ -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.
Expand Down
9 changes: 9 additions & 0 deletions tests/cmdline/commands/test_profile.py
Expand Up @@ -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
13 changes: 12 additions & 1 deletion tests/manage/configuration/test_configuration.py
Expand Up @@ -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):
Expand Down
4 changes: 2 additions & 2 deletions tests/storage/sqlite/test_archive.py
Expand Up @@ -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
Expand All @@ -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)

Expand Down

0 comments on commit 3510211

Please sign in to comment.