init: build broker url from its parts if available

Author: max-moser

CHANGES.rst

@@ -9,6 +9,10 @@
 Changes
 =======
 
+Version <next>
+
+- init: create broker url from its parts if available
+
 Version 2.0.0 (released 2024-12-02)
 
 - setup: bump invenio dependencies

invenio_celery/ext.py

@@ -68,8 +68,41 @@ def load_entry_points(self):
                         packages, related_name=related_name, force=True
                     )
 
+    def build_broker_url(self, app):
+        """Return the broker connection string if configured or build it from its parts.
+
+        If set, then ``BROKER_URL`` will be returned.
+        Otherwise, the URI will be pieced together by the configuration items
+        ``AMQP_BROKER_{USER,PASSWORD,HOST,PORT,NAME,PROTOCOL}``.
+        If that cannot be done (e.g. because required values are missing), then
+        ``None`` will be returned.
+
+        Note: see: https://docs.celeryq.dev/en/stable/userguide/configuration.html#new-lowercase-settings
+        """
+        if uri := app.config.get("BROKER_URL"):
+            return uri
+
+        params = {}
+        for config_name in ["USER", "PASSWORD", "HOST", "PORT", "VHOST", "PROTOCOL"]:
+            params[config_name] = app.config.get(f"AMQP_BROKER_{config_name}", None)
+
+        required_params = ["USER", "PASSWORD", "HOST", "PORT", "PROTOCOL"]
+        if all({params.get(p, None) for p in required_params}):
+            vhost = (params.get("VHOST") or "").lstrip("/")
+            uri = f"{params['PROTOCOL']}://{params['USER']}:{params['PASSWORD']}@{params['HOST']}:{params['PORT']}/{vhost}"
+            return uri
+        elif any(params.values()):
+            app.logger.warn(
+                'Ignoring "AMQP_BROKER_*" config values as they are only partially set.'
+            )
+
+        return None
+
     def init_config(self, app):
         """Initialize configuration."""
+        if broker_url := self.build_broker_url(app):
+            app.config["BROKER_URL"] = broker_url
+
         for k in dir(config):
             if k.startswith("CELERY_") or k.startswith("BROKER_"):
                 app.config.setdefault(k, getattr(config, k))

tests/test_invenio_celery.py

@@ -11,6 +11,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from flask import Flask
 from pkg_resources import EntryPoint
 
 from invenio_celery import InvenioCelery
@@ -147,3 +148,64 @@ def _mock():
 
     ext.get_active_tasks = _mock
     ext.suspend_queues(["feed"], sleep_time=0.1)
+
+
+@pytest.mark.parametrize(
+    "configs, expected_url",
+    [
+        (
+            {
+                "AMQP_BROKER_USER": "testuser",
+                "AMQP_BROKER_PASSWORD": "testpassword",
+                "AMQP_BROKER_HOST": "testhost",
+                "AMQP_BROKER_PORT": "5672",
+                "AMQP_BROKER_PROTOCOL": "amqp",
+                "AMQP_BROKER_VHOST": "/testvhost",
+            },
+            "amqp://testuser:testpassword@testhost:5672/testvhost",
+        ),
+        (
+            {
+                "AMQP_BROKER_USER": "testuser",
+                "AMQP_BROKER_PASSWORD": "testpassword",
+                "AMQP_BROKER_HOST": "testhost",
+                "AMQP_BROKER_PORT": "5672",
+                "AMQP_BROKER_PROTOCOL": "amqp",
+                "AMQP_BROKER_VHOST": "testvhost",
+            },
+            "amqp://testuser:testpassword@testhost:5672/testvhost",
+        ),
+        (
+            {
+                "AMQP_BROKER_USER": "testuser",
+                "AMQP_BROKER_PASSWORD": "testpassword",
+                "AMQP_BROKER_HOST": "testhost",
+                "AMQP_BROKER_PORT": "5672",
+                "AMQP_BROKER_PROTOCOL": "amqp",
+                "AMQP_BROKER_VHOST": "",
+            },
+            "amqp://testuser:testpassword@testhost:5672/",
+        ),
+        (
+            {
+                "AMQP_BROKER_USER": "testuser",
+                "AMQP_BROKER_PASSWORD": "testpassword",
+                "AMQP_BROKER_HOST": "testhost",
+                "AMQP_BROKER_PORT": "5672",
+                "AMQP_BROKER_PROTOCOL": "amqp",
+            },
+            "amqp://testuser:testpassword@testhost:5672/",
+        ),
+        (
+            {},
+            "redis://localhost:6379/0",
+        ),
+    ],
+)
+def test_build_broker_url_with_vhost(configs, expected_url):
+    """Test building broker URL with vhost."""
+    app = Flask("test_app")
+    assert "BROKER_URL" not in app.config
+    app.config.update(configs)
+    InvenioCelery(app)
+    assert app.config["BROKER_URL"] == expected_url