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/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/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" 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, + )