From 9fd595aa2176f32b336e05ea6f8c2d8a949a5ff7 Mon Sep 17 00:00:00 2001 From: Nicholas Devenish Date: Mon, 26 Oct 2026 13:14:23 +0000 Subject: [PATCH 1/2] Add CommonService.subscribe with overridable queue names Services can call self.subscribe to subscribe to a queue of recipe messages, with acknowledgement and recipe ID log extension enabled by default. The queue name can be overridden from the service's environment dictionary, which the service starter fills from a new --queue [NAME=]QUEUE option and the WORKFLOWS_QUEUE environment variable. Command line options take precedence over the variable. The option is only added if the starter's subclass has not already defined -q or --queue, so existing subclasses keep working. --- src/workflows/contrib/start_service.py | 64 ++++++++++ src/workflows/services/common_service.py | 68 ++++++++++ tests/contrib/test_start_service.py | 150 +++++++++++++++++++++++ tests/services/test_common_service.py | 45 +++++++ 4 files changed, 327 insertions(+) diff --git a/src/workflows/contrib/start_service.py b/src/workflows/contrib/start_service.py index cb6da236..1d680311 100644 --- a/src/workflows/contrib/start_service.py +++ b/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, diff --git a/src/workflows/services/common_service.py b/src/workflows/services/common_service.py index 4d6b9573..acb3ad9c 100644 --- a/src/workflows/services/common_service.py +++ b/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: diff --git a/tests/contrib/test_start_service.py b/tests/contrib/test_start_service.py index e7d6ec96..10fb0b24 100644 --- a/tests/contrib/test_start_service.py +++ b/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"}, + } diff --git a/tests/services/test_common_service.py b/tests/services/test_common_service.py index 91427a97..2f922c1d 100644 --- a/tests/services/test_common_service.py +++ b/tests/services/test_common_service.py @@ -626,3 +626,48 @@ def test_transport_callbacks_are_processed_from_main_queue(): service.SERVICE_STATUS_SHUTDOWN, service.SERVICE_STATUS_END, ] + + +@pytest.fixture +def wrap_subscribe(mocker): + return mocker.patch("workflows.recipe.wrap_subscribe") + + +def test_subscribe_uses_default_queue(wrap_subscribe): + service = CommonService() + service.transport = mock.Mock() + callback = mock.Mock() + service.subscribe("some_queue", callback, exclusive=True) + wrap_subscribe.assert_called_once_with( + service.transport, + "some_queue", + callback, + acknowledgement=True, + log_extender=service.extend_log, + exclusive=True, + ) + + +def test_subscribe_with_named_overrides(wrap_subscribe): + service = CommonService(environment={"queues": {"a": "a.test"}}) + service.transport = mock.Mock() + service.subscribe("a", mock.Mock()) + service.subscribe("b", mock.Mock(), acknowledgement=False) + assert [c.args[1] for c in wrap_subscribe.call_args_list] == ["a.test", "b"] + assert wrap_subscribe.call_args.kwargs["acknowledgement"] is False + + +def test_subscribe_bare_override_applies_to_single_queue_only(wrap_subscribe): + service = CommonService(environment={"queue": "x.test"}) + service.transport = mock.Mock() + service.subscribe("a", mock.Mock()) + assert wrap_subscribe.call_args.args[1] == "x.test" + with pytest.raises(RuntimeError, match="ambiguous"): + service.subscribe("b", mock.Mock()) + + +def test_subscribe_named_override_takes_precedence(wrap_subscribe): + service = CommonService(environment={"queue": "x.test", "queues": {"a": "a.test"}}) + service.transport = mock.Mock() + service.subscribe("a", mock.Mock()) + assert wrap_subscribe.call_args.args[1] == "a.test" From e1d05fc58e3b400e9182d4221646a62605d0c930 Mon Sep 17 00:00:00 2001 From: Nicholas Devenish Date: Mon, 26 Oct 2026 14:08:35 +0000 Subject: [PATCH 2/2] Use self.subscribe in the sample pipethrough service The sample consumer and transaction services subscribe to plain messages rather than recipes, so keep using transport.subscribe. --- src/workflows/services/sample_pipethrough.py | 7 +------ tests/services/test_sample_pipethrough.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 6 deletions(-) create mode 100644 tests/services/test_sample_pipethrough.py diff --git a/src/workflows/services/sample_pipethrough.py b/src/workflows/services/sample_pipethrough.py index e9975174..55543f9c 100644 --- a/src/workflows/services/sample_pipethrough.py +++ b/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.""" diff --git a/tests/services/test_sample_pipethrough.py b/tests/services/test_sample_pipethrough.py new file mode 100644 index 00000000..5c81680b --- /dev/null +++ b/tests/services/test_sample_pipethrough.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from unittest import mock + +from workflows.services.sample_pipethrough import SamplePipethrough + + +def test_service_subscribes_to_recipe_channel(mocker): + wrap_subscribe = mocker.patch("workflows.recipe.wrap_subscribe") + service = SamplePipethrough() + service.transport = mock.Mock() + service.initializing() + wrap_subscribe.assert_called_once_with( + service.transport, + "transient.destination", + service.process, + acknowledgement=False, + log_extender=service.extend_log, + )