Add CommonService.subscribe, with overridable queue names

src/workflows/contrib/start_service.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import optparse
+import os
 import sys
 from collections.abc import Callable
 from optparse import SUPPRESS_HELP, OptionParser
@@ -12,6 +13,31 @@
 import workflows.transport
 from workflows.transport.common_transport import CommonTransport
 
+# Environment variable holding queue overrides, as entries separated by commas
+# or whitespace. Each entry is NAME=QUEUE, to consume QUEUE in place of the
+# default queue NAME, or a bare QUEUE for a service consuming a single queue,
+# which may be given at most once. This is the same format as the --queue option.
+QUEUE_ENVIRONMENT_VARIABLE = "WORKFLOWS_QUEUE"
+
+
+def _parse_queue_overrides(
+    entries: list[str], source: str
+) -> tuple[str | None, dict[str, str]]:
+    """Split [NAME=]QUEUE entries into the unnamed queue and named overrides"""
+    queue: str | None = None
+    queues: dict[str, str] = {}
+    for entry in entries:
+        name, _, queue_name = entry.rpartition("=")
+        if not queue_name:
+            raise ValueError(f"Missing queue name in {source} entry {entry!r}")
+        if name:
+            queues[name] = queue_name
+        elif queue is not None:
+            raise ValueError(f"Only one {source} entry may be given without a NAME")
+        else:
+            queue = queue_name
+    return queue, queues
+
 
 class ServiceStarter:
     """A helper class to start a workflows service from the command line.
@@ -137,6 +163,25 @@ def run(
         # Call on_parser_preparation hook
         parser = self.on_parser_preparation(parser) or parser
 
+        # Add the queue option after the hook, and only if a subclass has not
+        # already defined its own -q or --queue option
+        if not (parser.has_option("-q") or parser.has_option("--queue")):
+            parser.add_option(
+                "-q",
+                "--queue",
+                action="append",
+                dest="workflows_queue",
+                default=[],
+                metavar="[NAME=]QUEUE",
+                help=(
+                    "Consume messages from QUEUE instead of the default queue NAME. "
+                    "NAME may be omitted for services that consume a single queue. "
+                    "May be specified multiple times. Defaults can be set in the "
+                    f"{QUEUE_ENVIRONMENT_VARIABLE} environment variable, as a comma or "
+                    "space separated list."
+                ),
+            )
+
         # Parse command line options
         (options, args) = parser.parse_args(cmdline_args)
 
@@ -208,6 +253,25 @@ def on_transport_preparation_hook() -> CommonTransport:
         if add_metrics_option and options.metrics:
             kwargs["environment"]["metrics"] = {"port": options.metrics_port}
 
+        # Queue overrides from the command line take precedence over those
+        # from the environment variable
+        try:
+            env_queue, env_queues = _parse_queue_overrides(
+                os.environ.get(QUEUE_ENVIRONMENT_VARIABLE, "")
+                .replace(",", " ")
+                .split(),
+                QUEUE_ENVIRONMENT_VARIABLE,
+            )
+            cli_queue, cli_queues = _parse_queue_overrides(
+                getattr(options, "workflows_queue", []), "--queue"
+            )
+        except ValueError as e:
+            parser.error(str(e))
+        if queue := cli_queue or env_queue:
+            kwargs["environment"]["queue"] = queue
+        if queues := {**env_queues, **cli_queues}:
+            kwargs["environment"]["queues"] = queues
+
         if options.liveness:
             kwargs["environment"]["liveness"] = {
                 "port": options.liveness_port,

src/workflows/services/common_service.py

@@ -19,6 +19,7 @@
 
 import workflows
 import workflows.logging
+import workflows.recipe
 from workflows.transport.common_transport import CommonTransport, MessageCallback
 from workflows.transport.middleware.otel_tracing import OTELTracingMiddleware
 
@@ -168,6 +169,7 @@ def __init__(self, *, environment: dict[str, Any] | None = None):
         )
         self._idle_callback: Callable[[], None] | None = None
         self._idle_time: float | None = None
+        self.__subscribed_channels: list[str] = []
         self.start_kwargs = {}
 
         # Logger will be overwritten in start() function
@@ -406,6 +408,72 @@ def _register_idle(self, idle_time: float, callback: Callable[[], None]) -> None
         self._idle_callback = callback
         self._idle_time = idle_time
 
+    def _resolve_queue(self, channel: str) -> str:
+        """Return the name of the queue to subscribe to in place of channel.
+
+        Overrides are read from the ``environment`` dictionary passed to the
+        service constructor (``self._environment``), not from operating system
+        environment variables or the recipe environment. The service starter
+        fills these entries from its --queue option and the WORKFLOWS_QUEUE
+        environment variable:
+
+        - ``environment["queues"]``, a mapping of default queue names to the
+          queue names to use instead
+        - ``environment["queue"]``, a single queue name. This is only accepted
+          for services subscribing to a single queue.
+        """
+        if override := (self._environment.get("queues") or {}).get(channel):
+            return override
+        if override := self._environment.get("queue"):
+            if self.__subscribed_channels:
+                raise RuntimeError(
+                    f"Queue override {override!r} is ambiguous as this service "
+                    "subscribes to multiple queues. Use --queue NAME=QUEUE instead."
+                )
+            return override
+        return channel
+
+    def subscribe(
+        self,
+        channel: str,
+        callback: Callable[..., Any],
+        *,
+        acknowledgement: bool = True,
+        **kwargs: Any,
+    ) -> int:
+        """Subscribe to a queue of recipe messages.
+
+        Messages are unpacked with :func:`workflows.recipe.wrap_subscribe`, so
+        the callback receives ``(recipe_wrapper, header, message)``. Use
+        ``self.transport.subscribe`` for a subscription without recipe handling.
+
+        The queue name can be overridden when the service is started, with the
+        --queue option or the WORKFLOWS_QUEUE environment variable, see
+        :meth:`_resolve_queue`. By default messages must be acknowledged, and
+        log messages are extended with the recipe ID.
+
+        Args:
+            channel: Default name of the queue to subscribe to.
+            callback: Function to be called when messages are received.
+            acknowledgement: Whether received messages must be acknowledged.
+            **kwargs: Passed on to :func:`workflows.recipe.wrap_subscribe`.
+
+        Returns:
+            A unique subscription ID
+        """
+        queue_name = self._resolve_queue(channel)
+        if queue_name != channel:
+            self.log.info("Subscribing to queue %s in place of %s", queue_name, channel)
+        self.__subscribed_channels.append(channel)
+        kwargs.setdefault("log_extender", self.extend_log)
+        return workflows.recipe.wrap_subscribe(
+            self.transport,
+            queue_name,
+            callback,
+            acknowledgement=acknowledgement,
+            **kwargs,
+        )
+
     def __update_service_status(self, statuscode: int) -> None:
         """Set the internal status of the service object, and notify frontend."""
         if self.__service_status != statuscode:

src/workflows/services/sample_pipethrough.py

@@ -4,7 +4,6 @@
 import time
 from typing import Any
 
-import workflows.recipe
 from workflows.recipe.wrapper import RecipeWrapper
 from workflows.services.common_service import CommonService
 
@@ -22,11 +21,7 @@ class SamplePipethrough(CommonService):
 
     def initializing(self) -> None:
         """Subscribe to a channel."""
-        workflows.recipe.wrap_subscribe(
-            self.transport,
-            "transient.destination",
-            self.process,
-        )
+        self.subscribe("transient.destination", self.process, acknowledgement=False)
 
     def process(self, rw: RecipeWrapper, header: dict, message: Any) -> None:
         """Consume message and send to output pipe."""

tests/contrib/test_start_service.py

@@ -30,6 +30,7 @@ def test_script_initialises_transport_and_starts_frontend(
     mock_options.service = "someservice"
     mock_options.transport = mock.sentinel.transport
     mock_options.liveness = False
+    mock_options.workflows_queue = []
     mock_parser.return_value.parse_args.return_value = (mock_options, mock.Mock())
     mock_services.get_known_services.return_value = {"SomeService": None}
 
@@ -59,6 +60,7 @@ def test_add_metrics_option(mock_services, mock_frontend, mock_tlookup, mock_par
     mock_options.service = "someservice"
     mock_options.transport = mock.sentinel.transport
     mock_options.liveness = False
+    mock_options.workflows_queue = []
     mock_options.metrics = False
     mock_options.metrics_port = 4242
     mock_parser.return_value.parse_args.return_value = (mock_options, mock.Mock())
@@ -89,3 +91,151 @@ def test_add_metrics_option(mock_services, mock_frontend, mock_tlookup, mock_par
         transport=mock.ANY,
         environment={"metrics": {"port": 4242}},
     )
+
+
+@pytest.mark.parametrize(
+    "args,environment",
+    [
+        ([], {}),
+        (["-q", "other"], {"queue": "other"}),
+        (
+            ["--queue", "a=a.test", "-q", "b=b.test"],
+            {"queues": {"a": "a.test", "b": "b.test"}},
+        ),
+        (
+            ["-q", "a=a.test", "-q", "other"],
+            {"queue": "other", "queues": {"a": "a.test"}},
+        ),
+    ],
+)
+@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
+@mock.patch("workflows.contrib.start_service.workflows.frontend")
+@mock.patch("workflows.contrib.start_service.workflows.services")
+def test_queue_option(mock_services, mock_frontend, mock_tlookup, args, environment):
+    mock_services.get_known_services.return_value = {"SomeService": None}
+    workflows.contrib.start_service.ServiceStarter().run(
+        cmdline_args=["-s", "SomeService", *args]
+    )
+    assert mock_frontend.Frontend.call_args.kwargs["environment"] == environment
+
+
+@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
+@mock.patch("workflows.contrib.start_service.workflows.frontend")
+@mock.patch("workflows.contrib.start_service.workflows.services")
+def test_queue_option_allows_one_unnamed_queue(
+    mock_services, mock_frontend, mock_tlookup, capsys
+):
+    mock_services.get_known_services.return_value = {"SomeService": None}
+    with pytest.raises(SystemExit):
+        workflows.contrib.start_service.ServiceStarter().run(
+            cmdline_args=["-s", "SomeService", "-q", "a", "-q", "b"]
+        )
+    assert "Only one --queue entry" in capsys.readouterr().err
+    mock_frontend.Frontend.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "variable,args,environment",
+    [
+        ("other", [], {"queue": "other"}),
+        ("a=a.env, b=b.env", [], {"queues": {"a": "a.env", "b": "b.env"}}),
+        (
+            "a=a.env b=b.env other",
+            [],
+            {"queue": "other", "queues": {"a": "a.env", "b": "b.env"}},
+        ),
+        # Command line options override the environment variable
+        (
+            "a=a.env,b=b.env",
+            ["-q", "a=a.cli"],
+            {"queues": {"a": "a.cli", "b": "b.env"}},
+        ),
+        ("env", ["-q", "cli"], {"queue": "cli"}),
+        ("env", ["-q", "a=a.cli"], {"queue": "env", "queues": {"a": "a.cli"}}),
+        ("", [], {}),
+    ],
+)
+@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
+@mock.patch("workflows.contrib.start_service.workflows.frontend")
+@mock.patch("workflows.contrib.start_service.workflows.services")
+def test_queue_environment_variable(
+    mock_services, mock_frontend, mock_tlookup, monkeypatch, variable, args, environment
+):
+    monkeypatch.setenv("WORKFLOWS_QUEUE", variable)
+    mock_services.get_known_services.return_value = {"SomeService": None}
+    workflows.contrib.start_service.ServiceStarter().run(
+        cmdline_args=["-s", "SomeService", *args]
+    )
+    assert mock_frontend.Frontend.call_args.kwargs["environment"] == environment
+
+
+@pytest.mark.parametrize(
+    "variable,args,error",
+    [
+        ("a b", [], "Only one WORKFLOWS_QUEUE entry"),
+        ("a=", [], "Missing queue name in WORKFLOWS_QUEUE entry 'a='"),
+        ("", ["-q", "a="], "Missing queue name in --queue entry 'a='"),
+    ],
+)
+@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
+@mock.patch("workflows.contrib.start_service.workflows.frontend")
+@mock.patch("workflows.contrib.start_service.workflows.services")
+def test_queue_overrides_rejected(
+    mock_services,
+    mock_frontend,
+    mock_tlookup,
+    monkeypatch,
+    capsys,
+    variable,
+    args,
+    error,
+):
+    monkeypatch.setenv("WORKFLOWS_QUEUE", variable)
+    mock_services.get_known_services.return_value = {"SomeService": None}
+    with pytest.raises(SystemExit):
+        workflows.contrib.start_service.ServiceStarter().run(
+            cmdline_args=["-s", "SomeService", *args]
+        )
+    assert error in capsys.readouterr().err
+    mock_frontend.Frontend.assert_not_called()
+
+
+class StarterWithOwnQueueOption(workflows.contrib.start_service.ServiceStarter):
+    """A subclass defining its own queue option, as dlstbx did before workflows"""
+
+    def __init__(self, *option_strings: str):
+        self.option_strings = option_strings or ("-q", "--queue")
+
+    def on_parser_preparation(self, parser):
+        parser.add_option(*self.option_strings, dest="queue", default=None)
+
+    def on_parsing(self, options, args):
+        self.options = options
+
+    def before_frontend_construction(self, kwargs):
+        if self.options.queue:
+            kwargs["environment"]["queue"] = self.options.queue
+        return kwargs
+
+
+@pytest.mark.parametrize("option_strings", [(), ("--queue",), ("-q", "--quiet")])
+@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
+@mock.patch("workflows.contrib.start_service.workflows.frontend")
+@mock.patch("workflows.contrib.start_service.workflows.services")
+def test_queue_option_skipped_if_already_defined(
+    mock_services, mock_frontend, mock_tlookup, monkeypatch, option_strings
+):
+    monkeypatch.setenv("WORKFLOWS_QUEUE", "a=a.env")
+    mock_services.get_known_services.return_value = {"SomeService": None}
+    starter = StarterWithOwnQueueOption(*option_strings)
+    flag = (option_strings or ("-q",))[0]
+    starter.run(cmdline_args=["-s", "SomeService", flag, "own"])
+
+    # The subclass option is parsed as before, as a single value
+    assert starter.options.queue == "own"
+    assert not hasattr(starter.options, "workflows_queue")
+    # The environment variable still applies
+    assert mock_frontend.Frontend.call_args.kwargs["environment"] == {
+        "queue": "own",
+        "queues": {"a": "a.env"},
+    }