From 7f03ed83fdd7465acc10312a121d328b56219261 Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Thu, 30 May 2024 21:00:17 +0200 Subject: [PATCH] CLI: Add automatic PostgreSQL management to `verdi presto` (#6432) The functionality of `verdi quicksetup` that automatically creates a PostgreSQL user and database was still relied upon by quite a number of users. Therefore, this functionality is integrated into `verdi presto` so that `verdi quicksetup` can finally be deprecated. Since the crux of `verdi presto` is to be as automatic as possible and require no prompting and as little configuration as necessary, only a few options are exposed. To use PostgreSQL instead of SQLite, the `--use-postgres` switch needs to be toggled. Only the hostname and port of the PostgreSQL server can be configured, as well as the username and password of an existing user with which to connect. --- .docker/aiida-core-base/Dockerfile | 1 - .../s6-assets/config-quick-setup.yaml | 14 -- .../s6-assets/init/aiida-prepare.sh | 12 +- .docker/aiida-core-with-services/Dockerfile | 1 - docs/source/reference/command_line.rst | 51 ++++--- src/aiida/cmdline/commands/cmd_presto.py | 127 ++++++++++++++++-- src/aiida/cmdline/commands/cmd_setup.py | 7 +- tests/cmdline/commands/test_presto.py | 26 +++- 8 files changed, 182 insertions(+), 57 deletions(-) delete mode 100644 .docker/aiida-core-base/s6-assets/config-quick-setup.yaml diff --git a/.docker/aiida-core-base/Dockerfile b/.docker/aiida-core-base/Dockerfile index 7de1ddb5bd..4dd66eecfb 100644 --- a/.docker/aiida-core-base/Dockerfile +++ b/.docker/aiida-core-base/Dockerfile @@ -163,7 +163,6 @@ RUN mkdir -p "${CONDA_DIR}/etc/conda/activate.d" && \ fix-permissions "${CONDA_DIR}" # COPY AiiDA profile configuration for profile setup init script -COPY --chown="${SYSTEM_UID}:${SYSTEM_GID}" s6-assets/config-quick-setup.yaml "/aiida/assets/config-quick-setup.yaml" COPY s6-assets/s6-rc.d /etc/s6-overlay/s6-rc.d COPY s6-assets/init /etc/init RUN mkdir /etc/init/run-before-daemon-start && \ diff --git a/.docker/aiida-core-base/s6-assets/config-quick-setup.yaml b/.docker/aiida-core-base/s6-assets/config-quick-setup.yaml deleted file mode 100644 index f255702326..0000000000 --- a/.docker/aiida-core-base/s6-assets/config-quick-setup.yaml +++ /dev/null @@ -1,14 +0,0 @@ -db_engine: postgresql_psycopg2 -db_backend: core.psql_dos -db_host: database -db_port: 5432 -su_db_username: postgres -su_db_password: password -su_db_name: template1 -db_name: aiida_db -db_username: aiida -db_password: password -broker_host: messaging -broker_port: 5672 -broker_username: guest -broker_password: guest diff --git a/.docker/aiida-core-base/s6-assets/init/aiida-prepare.sh b/.docker/aiida-core-base/s6-assets/init/aiida-prepare.sh index 45c2fe25be..f23eeca564 100755 --- a/.docker/aiida-core-base/s6-assets/init/aiida-prepare.sh +++ b/.docker/aiida-core-base/s6-assets/init/aiida-prepare.sh @@ -19,14 +19,10 @@ verdi config set warnings.development_version False if [[ ${SETUP_DEFAULT_AIIDA_PROFILE:-true} == true ]] && ! verdi profile show ${AIIDA_PROFILE_NAME} &> /dev/null; then # Create AiiDA profile. - verdi quicksetup \ - --non-interactive \ - --profile "${AIIDA_PROFILE_NAME:-default}" \ - --email "${AIIDA_USER_EMAIL:-aiida@localhost}" \ - --first-name "${AIIDA_USER_FIRST_NAME:-Giuseppe}" \ - --last-name "${AIIDA_USER_LAST_NAME:-Verdi}" \ - --institution "${AIIDA_USER_INSTITUTION:-Khedivial}" \ - --config "${AIIDA_CONFIG_FILE:-/aiida/assets/config-quick-setup.yaml}" + verdi presto \ + --profile "${AIIDA_PROFILE_NAME:-default}" \ + --email "${AIIDA_USER_EMAIL:-aiida@localhost}" \ + --use-postgres # Setup and configure local computer. computer_name=localhost diff --git a/.docker/aiida-core-with-services/Dockerfile b/.docker/aiida-core-with-services/Dockerfile index 896645d549..119cf7d0d1 100644 --- a/.docker/aiida-core-with-services/Dockerfile +++ b/.docker/aiida-core-with-services/Dockerfile @@ -33,7 +33,6 @@ RUN apt-get update --yes && \ fix-permissions /opt/rabbitmq_server-${RMQ_VERSION} # s6-overlay to start services -COPY --chown="${SYSTEM_UID}:${SYSTEM_GID}" s6-assets/config-quick-setup.yaml "/aiida/assets/config-quick-setup.yaml" COPY s6-assets/s6-rc.d /etc/s6-overlay/s6-rc.d COPY s6-assets/init /etc/init diff --git a/docs/source/reference/command_line.rst b/docs/source/reference/command_line.rst index 8c47f04046..bf4d668928 100644 --- a/docs/source/reference/command_line.rst +++ b/docs/source/reference/command_line.rst @@ -325,31 +325,50 @@ Below is a list with all available subcommands. Set up a new profile in a jiffy. - This command aims to make setting up a new profile as easy as possible. It intentionally - provides only a limited amount of options to customize the profile and by default does - not require any options to be specified at all. For full control, please use `verdi - profile setup`. + This command aims to make setting up a new profile as easy as possible. It does not + require any services, such as PostgreSQL and RabbitMQ. It intentionally provides only a + limited amount of options to customize the profile and by default does not require any + options to be specified at all. To create a new profile with full control over its + configuration, please use `verdi profile setup` instead. After running `verdi presto` you can immediately start using AiiDA without additional - setup. The created profile uses the `core.sqlite_dos` storage plugin which does not - require any services, such as PostgreSQL. The broker service RabbitMQ is also optional. - The command tries to connect to it using default settings and configures it for the - profile if found. Otherwise, the profile is created without a broker, in which case some - functionality will be unavailable, most notably running the daemon and submitting - processes to said daemon. - - The command performs the following actions: + setup. The command performs the following actions: * Create a new profile that is set as the new default * Create a default user for the profile (email can be configured through the `--email` option) * Set up the localhost as a `Computer` and configure it * Set a number of configuration options with sensible defaults + By default the command creates a profile that uses SQLite for the database. It + automatically checks for RabbitMQ running on the localhost, and, if it can connect, + configures that as the broker for the profile. Otherwise, the profile is created without + a broker, in which case some functionality will be unavailable, most notably running the + daemon and submitting processes to said daemon. + + When the `--use-postgres` flag is toggled, the command tries to connect to the + PostgreSQL server with connection paramaters taken from the `--postgres-hostname`, + `--postgres-port`, `--postgres-username` and `--postgres-password` options. It uses + these credentials to try and automatically create a user and database. If successful, + the newly created profile uses the new PostgreSQL database instead of SQLite. + Options: - --profile-name TEXT Name of the profile. By default, a unique name starting with - `presto` is automatically generated. [default: (dynamic)] - --email TEXT Email of the default user. [default: aiida@localhost] - --help Show this message and exit. + --profile-name TEXT Name of the profile. By default, a unique name starting with + `presto` is automatically generated. [default: (dynamic)] + --email TEXT Email of the default user. [default: (dynamic)] + --use-postgres When toggled on, the profile uses a PostgreSQL database + instead of an SQLite one. The connection details to the + PostgreSQL server can be configured with the relevant options. + The command attempts to automatically create a user and + database to use for the profile, but this can fail depending + on the configuration of the server. + --postgres-hostname TEXT The hostname of the PostgreSQL server. + --postgres-port INTEGER The port of the PostgreSQL server. + --postgres-username TEXT The username of the PostgreSQL user that is authorized to + create new databases. + --postgres-password TEXT The password of the PostgreSQL user that is authorized to + create new databases. + -n, --non-interactive Never prompt, such as for sudo password. + --help Show this message and exit. .. _reference:command-line:verdi-process: diff --git a/src/aiida/cmdline/commands/cmd_presto.py b/src/aiida/cmdline/commands/cmd_presto.py index a887a1a42b..a02250aaca 100644 --- a/src/aiida/cmdline/commands/cmd_presto.py +++ b/src/aiida/cmdline/commands/cmd_presto.py @@ -17,6 +17,7 @@ import click from aiida.cmdline.commands.cmd_verdi import verdi +from aiida.cmdline.params import options from aiida.cmdline.utils import echo from aiida.manage.configuration import get_config_option @@ -41,6 +42,52 @@ def get_default_presto_profile_name(): return f'{DEFAULT_PROFILE_NAME_PREFIX}-{last_index + 1}' +def detect_postgres_config( + profile_name: str, + postgres_hostname: str, + postgres_port: int, + postgres_username: str, + postgres_password: str, + non_interactive: bool, +) -> dict[str, t.Any]: + """.""" + import secrets + + from aiida.manage.configuration.settings import AIIDA_CONFIG_FOLDER + from aiida.manage.external.postgres import Postgres + + dbinfo = { + 'host': postgres_hostname, + 'port': postgres_port, + 'user': postgres_username, + 'password': postgres_password, + } + postgres = Postgres(interactive=not non_interactive, quiet=False, dbinfo=dbinfo) + + if not postgres.is_connected: + echo.echo_critical(f'Failed to connect to the PostgreSQL server using parameters: {dbinfo}') + + database_name = f'aiida-{profile_name}' + database_username = f'aiida-{profile_name}' + database_password = secrets.token_hex(15) + + try: + database_username, database_name = postgres.create_dbuser_db_safe( + dbname=database_name, dbuser=database_username, dbpass=database_password + ) + except Exception as exception: + echo.echo_critical(f'Unable to automatically create the PostgreSQL user and database: {exception}') + + return { + 'database_host': postgres_hostname, + 'database_port': postgres_port, + 'database_name': database_name, + 'database_username': database_username, + 'database_password': database_password, + 'repository_uri': f'file://{AIIDA_CONFIG_FOLDER / "repository" / profile_name}', + } + + @verdi.command('presto') @click.option( '--profile-name', @@ -51,25 +98,53 @@ def get_default_presto_profile_name(): ) @click.option( '--email', - default=get_config_option('autofill.user.email') or 'aiida@localhost', + default=lambda: get_config_option('autofill.user.email') or 'aiida@localhost', show_default=True, help='Email of the default user.', ) +@click.option( + '--use-postgres', + is_flag=True, + help='When toggled on, the profile uses a PostgreSQL database instead of an SQLite one. The connection details to ' + 'the PostgreSQL server can be configured with the relevant options. The command attempts to automatically create a ' + 'user and database to use for the profile, but this can fail depending on the configuration of the server.', +) +@click.option('--postgres-hostname', type=str, default='localhost', help='The hostname of the PostgreSQL server.') +@click.option('--postgres-port', type=int, default=5432, help='The port of the PostgreSQL server.') +@click.option( + '--postgres-username', + type=str, + default='postgres', + help='The username of the PostgreSQL user that is authorized to create new databases.', +) +@click.option( + '--postgres-password', + type=str, + required=False, + help='The password of the PostgreSQL user that is authorized to create new databases.', +) +@options.NON_INTERACTIVE(help='Never prompt, such as for sudo password.') @click.pass_context -def verdi_presto(ctx, profile_name, email): +def verdi_presto( + ctx, + profile_name, + email, + use_postgres, + postgres_hostname, + postgres_port, + postgres_username, + postgres_password, + non_interactive, +): """Set up a new profile in a jiffy. - This command aims to make setting up a new profile as easy as possible. It intentionally provides only a limited - amount of options to customize the profile and by default does not require any options to be specified at all. For - full control, please use `verdi profile setup`. - - After running `verdi presto` you can immediately start using AiiDA without additional setup. The created profile - uses the `core.sqlite_dos` storage plugin which does not require any services, such as PostgreSQL. The broker - service RabbitMQ is also optional. The command tries to connect to it using default settings and configures it for - the profile if found. Otherwise, the profile is created without a broker, in which case some functionality will be - unavailable, most notably running the daemon and submitting processes to said daemon. + This command aims to make setting up a new profile as easy as possible. It does not require any services, such as + PostgreSQL and RabbitMQ. It intentionally provides only a limited amount of options to customize the profile and by + default does not require any options to be specified at all. To create a new profile with full control over its + configuration, please use `verdi profile setup` instead. - The command performs the following actions: + After running `verdi presto` you can immediately start using AiiDA without additional setup. The command performs + the following actions: \b * Create a new profile that is set as the new default @@ -77,14 +152,38 @@ def verdi_presto(ctx, profile_name, email): * Set up the localhost as a `Computer` and configure it * Set a number of configuration options with sensible defaults + By default the command creates a profile that uses SQLite for the database. It automatically checks for RabbitMQ + running on the localhost, and, if it can connect, configures that as the broker for the profile. Otherwise, the + profile is created without a broker, in which case some functionality will be unavailable, most notably running the + daemon and submitting processes to said daemon. + + When the `--use-postgres` flag is toggled, the command tries to connect to the PostgreSQL server with connection + paramaters taken from the `--postgres-hostname`, `--postgres-port`, `--postgres-username` and `--postgres-password` + options. It uses these credentials to try and automatically create a user and database. If successful, the newly + created profile uses the new PostgreSQL database instead of SQLite. """ from aiida.brokers.rabbitmq.defaults import detect_rabbitmq_config from aiida.common import exceptions from aiida.manage.configuration import create_profile, load_profile from aiida.orm import Computer - storage_config: dict[str, t.Any] = {} - storage_backend = 'core.sqlite_dos' + postgres_config_kwargs = { + 'profile_name': profile_name, + 'postgres_hostname': postgres_hostname, + 'postgres_port': postgres_port, + 'postgres_username': postgres_username, + 'postgres_password': postgres_password, + 'non_interactive': non_interactive, + } + storage_config: dict[str, t.Any] = detect_postgres_config(**postgres_config_kwargs) if use_postgres else {} + storage_backend = 'core.psql_dos' if storage_config else 'core.sqlite_dos' + + if use_postgres: + echo.echo_report( + '`--use-postgres` enabled and database creation successful: configuring the profile to use PostgreSQL.' + ) + else: + echo.echo_report('Option `--use-postgres` not enabled: configuring the profile to use SQLite.') broker_config = detect_rabbitmq_config() broker_backend = 'core.rabbitmq' if broker_config is not None else None diff --git a/src/aiida/cmdline/commands/cmd_setup.py b/src/aiida/cmdline/commands/cmd_setup.py index 1072db6126..2f76d56434 100644 --- a/src/aiida/cmdline/commands/cmd_setup.py +++ b/src/aiida/cmdline/commands/cmd_setup.py @@ -13,11 +13,12 @@ from aiida.cmdline.commands.cmd_verdi import verdi from aiida.cmdline.params import options from aiida.cmdline.params.options.commands import setup as options_setup -from aiida.cmdline.utils import echo +from aiida.cmdline.utils import decorators, echo from aiida.manage.configuration import Profile, load_profile @verdi.command('setup') +@decorators.deprecated_command('This command is deprecated, use `verdi profile setup core.psql_dos` instead.') @options.NON_INTERACTIVE() @options_setup.SETUP_PROFILE() @options_setup.SETUP_USER_EMAIL() @@ -137,6 +138,10 @@ def setup( @verdi.command('quicksetup') +@decorators.deprecated_command( + 'This command is deprecated. For a fully automated alternative, use `verdi presto --use-postgres` instead. ' + 'For full control, use `verdi profile setup core.psql_dos`.' +) @options.NON_INTERACTIVE() # Cannot use `default` because that will fail validation of the `ProfileParamType` if the profile already exists and it # will be validated before the prompt to choose another. The `contextual_default` however, will not trigger the diff --git a/tests/cmdline/commands/test_presto.py b/tests/cmdline/commands/test_presto.py index 8d5bba2d77..c8503e0577 100644 --- a/tests/cmdline/commands/test_presto.py +++ b/tests/cmdline/commands/test_presto.py @@ -30,14 +30,14 @@ def get_profile_names(self): @pytest.mark.usefixtures('empty_config') @pytest.mark.parametrize('with_broker', (True, False)) def test_presto(run_cli_command, with_broker, monkeypatch): - """Test the ``verdi presto``.""" + """Test ``verdi presto`` with and without a broker present.""" from aiida.brokers.rabbitmq import defaults if not with_broker: # Patch the RabbitMQ detection function to pretend it could not find the service monkeypatch.setattr(defaults, 'detect_rabbitmq_config', lambda: None) - result = run_cli_command(verdi_presto) + result = run_cli_command(verdi_presto, ['--non-interactive']) assert 'Created new profile `presto`.' in result.output with profile_context('presto', allow_switch=True) as profile: @@ -48,3 +48,25 @@ def test_presto(run_cli_command, with_broker, monkeypatch): assert profile.process_control_backend == 'core.rabbitmq' else: assert profile.process_control_backend is None + + +@pytest.mark.usefixtures('empty_config') +def test_presto_use_postgres(run_cli_command, manager): + """Test the ``verdi presto`` with the ``--use-postgres`` flag.""" + result = run_cli_command(verdi_presto, ['--non-interactive', '--use-postgres']) + assert 'Created new profile `presto`.' in result.output + + with profile_context('presto', allow_switch=True) as profile: + assert profile.name == 'presto' + localhost = Computer.collection.get(label='localhost') + assert localhost.is_configured + assert profile.storage_backend == 'core.psql_dos' + assert manager.get_profile_storage() + + +@pytest.mark.usefixtures('empty_config') +def test_presto_use_postgres_fail(run_cli_command): + """Test the ``verdi presto`` with the ``-use-postgres`` flag specifying an incorrect option.""" + options = ['--non-interactive', '--use-postgres', '--postgres-port', str(5000)] + result = run_cli_command(verdi_presto, options, raises=True) + assert 'Failed to connect to the PostgreSQL server' in result.output