From d79ffc756ba359e181993e52d5b3a962dfb5e2e0 Mon Sep 17 00:00:00 2001
From: Sebastiaan Huber <mail@sphuber.net>
Date: Wed, 17 Apr 2024 16:59:15 +0200
Subject: [PATCH] Automatically try and detect a RabbitMQ server

The `verdi presto` command now checks whether it can connect to a
RabbitMQ instance on the localhost with default connection parameters.
If it is found, the profile is configured with RabbitMQ as a broker,
otehrwise the profile will simply not configure a broker.
---
 src/aiida/brokers/rabbitmq/defaults.py   | 21 +++++++++++++++++++++
 src/aiida/cmdline/commands/cmd_presto.py | 13 ++++++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/src/aiida/brokers/rabbitmq/defaults.py b/src/aiida/brokers/rabbitmq/defaults.py
index 6e3684ae0e..c1e4daa5b8 100644
--- a/src/aiida/brokers/rabbitmq/defaults.py
+++ b/src/aiida/brokers/rabbitmq/defaults.py
@@ -1,5 +1,9 @@
 """Defaults related to RabbitMQ."""
 
+from __future__ import annotations
+
+import typing as t
+
 from aiida.common.extendeddicts import AttributeDict
 
 __all__ = ('BROKER_DEFAULTS',)
@@ -19,3 +23,20 @@
         'heartbeat': 600,
     }
 )
+
+
+def detect_rabbitmq_config() -> 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 = dict(BROKER_DEFAULTS)
+
+    try:
+        connect(connection_params=connection_params)
+    except ConnectionError:
+        return None
+
+    return connection_params
diff --git a/src/aiida/cmdline/commands/cmd_presto.py b/src/aiida/cmdline/commands/cmd_presto.py
index d3fa1c335a..4177d64fc6 100644
--- a/src/aiida/cmdline/commands/cmd_presto.py
+++ b/src/aiida/cmdline/commands/cmd_presto.py
@@ -93,12 +93,21 @@ def verdi_presto(ctx, profile_name, email, first_name, last_name, institution):
     """
     import tempfile
 
+    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_backend = 'core.sqlite_dos'
     storage_config: dict[str, t.Any] = {}
+    storage_backend = 'core.sqlite_dos'
+
+    broker_config = detect_rabbitmq_config()
+    broker_backend = 'core.rabbitmq' if broker_config is not None else None
+
+    if broker_config is None:
+        echo.echo_report('RabbitMQ server not found: configuring the profile without a broker.')
+    else:
+        echo.echo_report('RabbitMQ server detected: configuring the profile with a broker.')
 
     try:
         profile = create_profile(
@@ -110,6 +119,8 @@ def verdi_presto(ctx, profile_name, email, first_name, last_name, institution):
             institution=institution,
             storage_backend=storage_backend,
             storage_config=storage_config,
+            broker_backend=broker_backend,
+            broker_config=broker_config,
         )
     except (ValueError, TypeError, exceptions.EntryPointError, exceptions.StorageMigrationError) as exception:
         echo.echo_critical(str(exception))