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