From 5e242808f1b35e044971d717f75c84d2a3b2a81d Mon Sep 17 00:00:00 2001
From: Sebastiaan Huber <mail@sphuber.net>
Date: Wed, 5 Jun 2024 17:42:55 +0200
Subject: [PATCH] CLI: Add the `verdi profile configure-rabbitmq`

Now that profiles can be created without defining a broker, a command
is needed that can add a RabbitMQ connection configuration. The new
command `verdi profile configure-rabbitmq` enables a broker for a
profile if it wasn't already, and allows configuring the connection
parameters.
---
 docs/source/reference/command_line.rst        | 11 +++++----
 src/aiida/cmdline/commands/cmd_profile.py     | 20 ++++++++++++++++
 .../cmdline/params/options/commands/setup.py  |  2 ++
 tests/cmdline/commands/test_profile.py        | 23 +++++++++++++++++++
 4 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/docs/source/reference/command_line.rst b/docs/source/reference/command_line.rst
index bf4d668928..9c95d9e656 100644
--- a/docs/source/reference/command_line.rst
+++ b/docs/source/reference/command_line.rst
@@ -414,11 +414,12 @@ Below is a list with all available subcommands.
       --help  Show this message and exit.
 
     Commands:
-      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.
+      configure-rabbitmq  Configure RabbitMQ for a profile.
+      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.
 
 
 .. _reference:command-line:verdi-quicksetup:
diff --git a/src/aiida/cmdline/commands/cmd_profile.py b/src/aiida/cmdline/commands/cmd_profile.py
index aa743028a1..8edc47f9e9 100644
--- a/src/aiida/cmdline/commands/cmd_profile.py
+++ b/src/aiida/cmdline/commands/cmd_profile.py
@@ -128,6 +128,26 @@ def profile_setup():
     """Set up a new profile."""
 
 
+@verdi_profile.command('configure-rabbitmq')  # type: ignore[arg-type]
+@arguments.PROFILE(default=defaults.get_default_profile)
+@setup.SETUP_BROKER_PROTOCOL()
+@setup.SETUP_BROKER_USERNAME()
+@setup.SETUP_BROKER_PASSWORD()
+@setup.SETUP_BROKER_HOST()
+@setup.SETUP_BROKER_PORT()
+@setup.SETUP_BROKER_VIRTUAL_HOST()
+@options.NON_INTERACTIVE()
+@click.pass_context
+def profile_configure_rabbitmq(ctx, profile, **kwargs):
+    """Configure RabbitMQ for a profile.
+
+    Enable RabbitMQ for a profile that was created without a broker, or reconfigure existing connection details.
+    """
+    profile.set_process_controller(name='core.rabbitmq', config=kwargs)
+    ctx.obj.config.update_profile(profile)
+    ctx.obj.config.store()
+
+
 @verdi_profile.command('list')
 def profile_list():
     """Display a list of all available profiles."""
diff --git a/src/aiida/cmdline/params/options/commands/setup.py b/src/aiida/cmdline/params/options/commands/setup.py
index bbd980c976..008f51b3a0 100644
--- a/src/aiida/cmdline/params/options/commands/setup.py
+++ b/src/aiida/cmdline/params/options/commands/setup.py
@@ -50,6 +50,8 @@ def get_profile_attribute_default(attribute_tuple, ctx):
     try:
         data = ctx.params['profile'].dictionary
         for part in parts:
+            if data is None:
+                return default
             data = data[part]
         return data
     except KeyError:
diff --git a/tests/cmdline/commands/test_profile.py b/tests/cmdline/commands/test_profile.py
index 65c715e86a..9d0d2ca2a5 100644
--- a/tests/cmdline/commands/test_profile.py
+++ b/tests/cmdline/commands/test_profile.py
@@ -255,3 +255,26 @@ def test_setup_no_use_rabbitmq(run_cli_command, isolated_config):
     profile = isolated_config.get_profile(profile_name)
     assert profile.process_control_backend is None
     assert profile.process_control_config == {}
+
+
+def test_configure_rabbitmq(run_cli_command, isolated_config):
+    """Test the ``verdi profile configure-rabbitmq`` command."""
+    profile_name = 'profile'
+
+    # First setup a profile without a broker configured
+    options = ['core.sqlite_dos', '-n', '--email', 'a@a', '--profile', profile_name, '--no-use-rabbitmq']
+    run_cli_command(cmd_profile.profile_setup, options, use_subprocess=False)
+    profile = isolated_config.get_profile(profile_name)
+    assert profile.process_control_backend is None
+    assert profile.process_control_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)
+    assert profile.process_control_backend == 'core.rabbitmq'
+
+    # 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'