From b39076b6703b0ede8ddf85073d537677c1ae6436 Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Fri, 28 Jun 2024 10:24:28 +0200 Subject: [PATCH] CLI: Remove the RabbitMQ options from `verdi profile setup` (#6480) For the vast majority of use cases, users will have a default setup for RabbitMQ and so the default configuration will be adequate and so they will not need the options in the command. On the flipside, showing the options by default can makes the command harder to use as users will take pause to think what value to pass. Since there is the `verdi profile configure-rabbitmq` command now that allows to configure or reconfigure the RabbitMQ connection parameters for an existing profile, it is fine to remove these options from the profile setup. Advanced users that need to customize the connection parameters can resort to that separate command. --- src/aiida/brokers/rabbitmq/defaults.py | 7 +- src/aiida/cmdline/commands/cmd_presto.py | 40 +++++++--- src/aiida/cmdline/commands/cmd_profile.py | 90 +++++++++++------------ tests/cmdline/commands/test_presto.py | 5 +- tests/cmdline/commands/test_profile.py | 2 +- 5 files changed, 78 insertions(+), 66 deletions(-) diff --git a/src/aiida/brokers/rabbitmq/defaults.py b/src/aiida/brokers/rabbitmq/defaults.py index b312897f73..21a15f1ad0 100644 --- a/src/aiida/brokers/rabbitmq/defaults.py +++ b/src/aiida/brokers/rabbitmq/defaults.py @@ -36,10 +36,10 @@ def detect_rabbitmq_config( host: str | None = None, port: int | None = None, virtual_host: str | None = None, - heartbeat: int | None = None, -) -> dict[str, t.Any] | None: +) -> dict[str, t.Any]: """Try to connect to a RabbitMQ server with the default connection parameters. + :raises ConnectionError: If the connection failed with the provided connection parameters :returns: The connection parameters if the RabbitMQ server was successfully connected to, or ``None`` otherwise. """ from kiwipy.rmq.threadcomms import connect @@ -51,7 +51,6 @@ def detect_rabbitmq_config( 'host': host or os.getenv('AIIDA_BROKER_HOST', BROKER_DEFAULTS['host']), 'port': port or int(os.getenv('AIIDA_BROKER_PORT', BROKER_DEFAULTS['port'])), 'virtual_host': virtual_host or os.getenv('AIIDA_BROKER_VIRTUAL_HOST', BROKER_DEFAULTS['virtual_host']), - 'heartbeat': heartbeat or int(os.getenv('AIIDA_BROKER_HEARTBEAT', BROKER_DEFAULTS['heartbeat'])), } LOGGER.info(f'Attempting to connect to RabbitMQ with parameters: {connection_params}') @@ -59,7 +58,7 @@ def detect_rabbitmq_config( try: connect(connection_params=connection_params) except ConnectionError: - return None + raise ConnectionError(f'Failed to connect with following connection parameters: {connection_params}') # The profile configuration expects the keys of the broker config to be prefixed with ``broker_``. return {f'broker_{key}': value for key, value in connection_params.items()} diff --git a/src/aiida/cmdline/commands/cmd_presto.py b/src/aiida/cmdline/commands/cmd_presto.py index 1893a6d461..d0835b8956 100644 --- a/src/aiida/cmdline/commands/cmd_presto.py +++ b/src/aiida/cmdline/commands/cmd_presto.py @@ -50,7 +50,13 @@ def detect_postgres_config( postgres_password: str, non_interactive: bool, ) -> dict[str, t.Any]: - """.""" + """Attempt to connect to the given PostgreSQL server and create a new user and database. + + :raises ConnectionError: If no connection could be established to the PostgreSQL server or a user and database + could not be created. + :returns: The connection configuration for the newly created user and database as can be used directly for the + storage configuration of the ``core.psql_dos`` storage plugin. + """ import secrets from aiida.manage.configuration.settings import AIIDA_CONFIG_FOLDER @@ -65,7 +71,7 @@ def detect_postgres_config( 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}') + raise ConnectionError(f'Failed to connect to the PostgreSQL server using parameters: {dbinfo}') database_name = f'aiida-{profile_name}' database_username = f'aiida-{profile_name}' @@ -76,7 +82,7 @@ def detect_postgres_config( 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}') + raise ConnectionError(f'Unable to automatically create the PostgreSQL user and database: {exception}') return { 'database_hostname': postgres_hostname, @@ -175,23 +181,33 @@ def verdi_presto( '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' + + storage_backend: str = 'core.sqlite_dos' + storage_config: dict[str, t.Any] = {} if use_postgres: - echo.echo_report( - '`--use-postgres` enabled and database creation successful: configuring the profile to use PostgreSQL.' - ) + try: + storage_config = detect_postgres_config(**postgres_config_kwargs) + except ConnectionError as exception: + echo.echo_critical(str(exception)) + else: + echo.echo_report( + '`--use-postgres` enabled and database creation successful: configuring the profile to use PostgreSQL.' + ) + storage_backend = 'core.psql_dos' 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 + broker_backend = None + broker_config = None - if broker_config is None: - echo.echo_report('RabbitMQ server not found: configuring the profile without a broker.') + try: + broker_config = detect_rabbitmq_config() + except ConnectionError as exception: + echo.echo_report(f'RabbitMQ server not found ({exception}): configuring the profile without a broker.') else: echo.echo_report('RabbitMQ server detected: configuring the profile with a broker.') + broker_backend = 'core.rabbitmq' try: profile = create_profile( diff --git a/src/aiida/cmdline/commands/cmd_profile.py b/src/aiida/cmdline/commands/cmd_profile.py index 047126a38c..5e89b72d70 100644 --- a/src/aiida/cmdline/commands/cmd_profile.py +++ b/src/aiida/cmdline/commands/cmd_profile.py @@ -27,7 +27,17 @@ def verdi_profile(): def command_create_profile( - ctx: click.Context, storage_cls, non_interactive: bool, profile: Profile, set_as_default: bool = True, **kwargs + ctx: click.Context, + storage_cls, + non_interactive: bool, + profile: Profile, + set_as_default: bool = True, + email: str | None = None, + first_name: str | None = None, + last_name: str | None = None, + institution: str | None = None, + use_rabbitmq: bool = True, + **kwargs, ): """Create a new profile, initialise its storage and create a default user. @@ -37,43 +47,44 @@ def command_create_profile( :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 set_as_default: Whether to set the created profile as the new default. + :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 use_rabbitmq: Whether to configure RabbitMQ as the broker. :param kwargs: Arguments to initialise instance of the selected storage implementation. """ + from aiida.brokers.rabbitmq.defaults import detect_rabbitmq_config from aiida.plugins.entry_point import get_entry_point_from_class - if not storage_cls.read_only and kwargs.get('email', None) is None: + if not storage_cls.read_only and email is None: raise click.BadParameter('The option is required for storages that are not read-only.', param_hint='--email') - email = kwargs.pop('email') - first_name = kwargs.pop('first_name') - last_name = kwargs.pop('last_name') - institution = kwargs.pop('institution') - _, storage_entry_point = get_entry_point_from_class(storage_cls.__module__, storage_cls.__name__) assert storage_entry_point is not None - if kwargs.pop('use_rabbitmq'): - broker_backend = 'core.rabbitmq' - broker_config = { - key: kwargs.get(key) - for key in ( - 'broker_protocol', - 'broker_username', - 'broker_password', - 'broker_host', - 'broker_port', - 'broker_virtual_host', - ) - } + broker_backend = None + broker_config = None + + if use_rabbitmq: + try: + broker_config = detect_rabbitmq_config() + except ConnectionError as exception: + echo.echo_warning(f'RabbitMQ server not reachable: {exception}.') + else: + echo.echo_success(f'RabbitMQ server detected with connection parameters: {broker_config}') + broker_backend = 'core.rabbitmq' + + echo.echo_report('RabbitMQ can be reconfigured with `verdi profile configure-rabbitmq`.') else: - broker_backend = None - broker_config = None + echo.echo_report('Creating profile without RabbitMQ.') + echo.echo_report('It can be configured at a later point in time with `verdi profile configure-rabbitmq`.') try: profile = create_profile( ctx.obj.config, name=profile.name, - email=email, + email=email, # type: ignore[arg-type] first_name=first_name, last_name=last_name, institution=institution, @@ -104,24 +115,6 @@ def command_create_profile( setup.SETUP_USER_LAST_NAME(), setup.SETUP_USER_INSTITUTION(), setup.SETUP_USE_RABBITMQ(), - setup.SETUP_BROKER_PROTOCOL( - prompt_fn=lambda ctx: ctx.params['use_rabbitmq'], required_fn=lambda ctx: ctx.params['use_rabbitmq'] - ), - setup.SETUP_BROKER_USERNAME( - prompt_fn=lambda ctx: ctx.params['use_rabbitmq'], required_fn=lambda ctx: ctx.params['use_rabbitmq'] - ), - setup.SETUP_BROKER_PASSWORD( - prompt_fn=lambda ctx: ctx.params['use_rabbitmq'], required_fn=lambda ctx: ctx.params['use_rabbitmq'] - ), - setup.SETUP_BROKER_HOST( - prompt_fn=lambda ctx: ctx.params['use_rabbitmq'], required_fn=lambda ctx: ctx.params['use_rabbitmq'] - ), - setup.SETUP_BROKER_PORT( - prompt_fn=lambda ctx: ctx.params['use_rabbitmq'], required_fn=lambda ctx: ctx.params['use_rabbitmq'] - ), - setup.SETUP_BROKER_VIRTUAL_HOST( - prompt_fn=lambda ctx: ctx.params['use_rabbitmq'], required_fn=lambda ctx: ctx.params['use_rabbitmq'] - ), ], ) def profile_setup(): @@ -146,22 +139,23 @@ def profile_configure_rabbitmq(ctx, profile, non_interactive, force, **kwargs): """ from aiida.brokers.rabbitmq.defaults import detect_rabbitmq_config - connection_params = {key.lstrip('broker_'): value for key, value in kwargs.items() if key.startswith('broker_')} - - broker_config = detect_rabbitmq_config(**connection_params) + broker_config = {key: value for key, value in kwargs.items() if key.startswith('broker_')} + connection_params = {key.lstrip('broker_'): value for key, value in broker_config.items()} - if broker_config is None: - echo.echo_warning(f'Unable to connect to RabbitMQ server with configuration: {connection_params}') + try: + broker_config = detect_rabbitmq_config(**connection_params) + except ConnectionError as exception: + echo.echo_warning(f'Unable to connect to RabbitMQ server: {exception}') if not force: click.confirm('Do you want to continue with the provided configuration?', abort=True) else: echo.echo_success('Connected to RabbitMQ with the provided connection parameters') - profile.set_process_controller(name='core.rabbitmq', config=kwargs) + profile.set_process_controller(name='core.rabbitmq', config=broker_config) ctx.obj.config.update_profile(profile) ctx.obj.config.store() - echo.echo_success(f'RabbitMQ configuration for `{profile.name}` updated to: {connection_params}') + echo.echo_success(f'RabbitMQ configuration for `{profile.name}` updated to: {broker_config}') @verdi_profile.command('list') diff --git a/tests/cmdline/commands/test_presto.py b/tests/cmdline/commands/test_presto.py index 8651664f61..3ec1d1e5da 100644 --- a/tests/cmdline/commands/test_presto.py +++ b/tests/cmdline/commands/test_presto.py @@ -32,8 +32,11 @@ def test_presto_without_rmq(pytestconfig, run_cli_command, monkeypatch): """Test the ``verdi presto`` without RabbitMQ.""" from aiida.brokers.rabbitmq import defaults + def detect_rabbitmq_config(**kwargs): + raise ConnectionError() + # Patch the RabbitMQ detection function to pretend it could not find the service - monkeypatch.setattr(defaults, 'detect_rabbitmq_config', lambda: None) + monkeypatch.setattr(defaults, 'detect_rabbitmq_config', lambda: detect_rabbitmq_config()) result = run_cli_command(verdi_presto, ['--non-interactive']) assert 'Created new profile `presto`.' in result.output diff --git a/tests/cmdline/commands/test_profile.py b/tests/cmdline/commands/test_profile.py index 781a8b3cfa..7f89cd2f31 100644 --- a/tests/cmdline/commands/test_profile.py +++ b/tests/cmdline/commands/test_profile.py @@ -299,7 +299,7 @@ def test_configure_rabbitmq(run_cli_command, isolated_config): # Verify that configuring with incorrect options and `--force` raises a warning but still configures the broker options = [profile_name, '-f', '--broker-port', '1234'] cli_result = run_cli_command(cmd_profile.profile_configure_rabbitmq, options, use_subprocess=False) - assert 'Unable to connect to RabbitMQ server with configuration:' in cli_result.stdout + assert 'Unable to connect to RabbitMQ server: Failed to connect' in cli_result.stdout assert profile.process_control_config['broker_port'] == 1234 # Call it again to check it works to reconfigure existing broker connection parameters