Skip to content

Commit

Permalink
👌 CLI: Give feedback for configure-rabbitmq
Browse files Browse the repository at this point in the history
Currently the `verdi profile configure-rabbitmq` command doesn't give any feedback to
the user whether the provided options can successfully connect to the RabbitMQ server.
Here we adapt the `detect_rabbitmq_config` function to accept the broker configuration
as `**kwargs`, and use it to check if the provided options in the `configure-rabbitmq`
can successfully connect to the RabbitMQ server. A "success" message is printed if
we can connect to the server, else a warning is printed and the user is asked for
confirmation before proceeding to configure the broker.

A `--force` flag is also added to avoid asking for confirmation in case the command
is unable to connect to the broker.
  • Loading branch information
mbercx authored and sphuber committed Jun 19, 2024
1 parent aaada54 commit c2ca642
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 14 deletions.
24 changes: 16 additions & 8 deletions src/aiida/brokers/rabbitmq/defaults.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,21 +29,29 @@
)


def detect_rabbitmq_config() -> dict[str, t.Any] | None:
def detect_rabbitmq_config(
protocol: str | None = None,
username: str | None = None,
password: str | None = None,
host: str | None = None,
port: int | None = None,
virtual_host: str | None = None,
heartbeat: int | None = None,
) -> dict[str, t.Any] | None:
"""Try to connect to a RabbitMQ server with the default connection parameters.
:returns: The connection parameters if the RabbitMQ server was successfully connected to, or ``None`` otherwise.
"""
from kiwipy.rmq.threadcomms import connect

connection_params = {
'protocol': os.getenv('AIIDA_BROKER_PROTOCOL', BROKER_DEFAULTS['protocol']),
'username': os.getenv('AIIDA_BROKER_USERNAME', BROKER_DEFAULTS['username']),
'password': os.getenv('AIIDA_BROKER_PASSWORD', BROKER_DEFAULTS['password']),
'host': os.getenv('AIIDA_BROKER_HOST', BROKER_DEFAULTS['host']),
'port': os.getenv('AIIDA_BROKER_PORT', BROKER_DEFAULTS['port']),
'virtual_host': os.getenv('AIIDA_BROKER_VIRTUAL_HOST', BROKER_DEFAULTS['virtual_host']),
'heartbeat': os.getenv('AIIDA_BROKER_HEARTBEAT', BROKER_DEFAULTS['heartbeat']),
'protocol': protocol or os.getenv('AIIDA_BROKER_PROTOCOL', BROKER_DEFAULTS['protocol']),
'username': username or os.getenv('AIIDA_BROKER_USERNAME', BROKER_DEFAULTS['username']),
'password': password or os.getenv('AIIDA_BROKER_PASSWORD', BROKER_DEFAULTS['password']),
'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}')
Expand Down
18 changes: 17 additions & 1 deletion src/aiida/cmdline/commands/cmd_profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ def profile_setup():

@verdi_profile.command('configure-rabbitmq') # type: ignore[arg-type]
@arguments.PROFILE(default=defaults.get_default_profile)
@options.FORCE()
@setup.SETUP_BROKER_PROTOCOL()
@setup.SETUP_BROKER_USERNAME()
@setup.SETUP_BROKER_PASSWORD()
Expand All @@ -138,15 +139,30 @@ def profile_setup():
@setup.SETUP_BROKER_VIRTUAL_HOST()
@options.NON_INTERACTIVE(default=True, show_default='--non-interactive')
@click.pass_context
def profile_configure_rabbitmq(ctx, profile, **kwargs):
def profile_configure_rabbitmq(ctx, profile, non_interactive, force, **kwargs):
"""Configure RabbitMQ for a profile.
Enable RabbitMQ for a profile that was created without a broker, or reconfigure existing connection details.
"""
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)

if broker_config is None:
echo.echo_warning(f'Unable to connect to RabbitMQ server with configuration: {connection_params}')
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)
ctx.obj.config.update_profile(profile)
ctx.obj.config.store()

echo.echo_success(f'RabbitMQ configuration for `{profile.name}` updated to: {connection_params}')


@verdi_profile.command('list')
def profile_list():
Expand Down
18 changes: 13 additions & 5 deletions tests/cmdline/commands/test_profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -284,18 +284,26 @@ def test_configure_rabbitmq(run_cli_command, isolated_config):

# Now run the command to configure the broker
options = [profile_name, '-n']
run_cli_command(cmd_profile.profile_configure_rabbitmq, options, use_subprocess=False)
cli_result = run_cli_command(cmd_profile.profile_configure_rabbitmq, options, use_subprocess=False)
assert profile.process_control_backend == 'core.rabbitmq'
assert 'Connected to RabbitMQ with the provided connection parameters' in cli_result.stdout

# Verify that running in non-interactive mode is the default
options = [
profile_name,
]
run_cli_command(cmd_profile.profile_configure_rabbitmq, options, use_subprocess=True)
assert profile.process_control_backend == 'core.rabbitmq'
assert 'Connected to RabbitMQ with the provided connection parameters' in cli_result.stdout

# 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 profile.process_control_config['broker_port'] == 1234

# Call it again to check it works to reconfigure existing broker connection parameters
options = [profile_name, '-n', '--broker-host', 'rabbitmq.broker.com']
run_cli_command(cmd_profile.profile_configure_rabbitmq, options, use_subprocess=False)
assert profile.process_control_backend == 'core.rabbitmq'
assert profile.process_control_config['broker_host'] == 'rabbitmq.broker.com'
options = [profile_name, '-n', '--broker-port', '5672']
cli_result = run_cli_command(cmd_profile.profile_configure_rabbitmq, options, use_subprocess=False)
assert 'Connected to RabbitMQ with the provided connection parameters' in cli_result.stdout
assert profile.process_control_config['broker_port'] == 5672

0 comments on commit c2ca642

Please sign in to comment.