Skip to content

Commit 9fd595a

Browse files
committed
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.
1 parent ee9c938 commit 9fd595a

4 files changed

Lines changed: 327 additions & 0 deletions

File tree

‎src/workflows/contrib/start_service.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from __future__ import annotations
22

33
import optparse
4+
import os
45
import sys
56
from collections.abc import Callable
67
from optparse import SUPPRESS_HELP, OptionParser
@@ -12,6 +13,31 @@
1213
import workflows.transport
1314
from workflows.transport.common_transport import CommonTransport
1415

16+
# Environment variable holding queue overrides, as entries separated by commas
17+
# or whitespace. Each entry is NAME=QUEUE, to consume QUEUE in place of the
18+
# default queue NAME, or a bare QUEUE for a service consuming a single queue,
19+
# which may be given at most once. This is the same format as the --queue option.
20+
QUEUE_ENVIRONMENT_VARIABLE = "WORKFLOWS_QUEUE"
21+
22+
23+
def _parse_queue_overrides(
24+
entries: list[str], source: str
25+
) -> tuple[str | None, dict[str, str]]:
26+
"""Split [NAME=]QUEUE entries into the unnamed queue and named overrides"""
27+
queue: str | None = None
28+
queues: dict[str, str] = {}
29+
for entry in entries:
30+
name, _, queue_name = entry.rpartition("=")
31+
if not queue_name:
32+
raise ValueError(f"Missing queue name in {source} entry {entry!r}")
33+
if name:
34+
queues[name] = queue_name
35+
elif queue is not None:
36+
raise ValueError(f"Only one {source} entry may be given without a NAME")
37+
else:
38+
queue = queue_name
39+
return queue, queues
40+
1541

1642
class ServiceStarter:
1743
"""A helper class to start a workflows service from the command line.
@@ -137,6 +163,25 @@ def run(
137163
# Call on_parser_preparation hook
138164
parser = self.on_parser_preparation(parser) or parser
139165

166+
# Add the queue option after the hook, and only if a subclass has not
167+
# already defined its own -q or --queue option
168+
if not (parser.has_option("-q") or parser.has_option("--queue")):
169+
parser.add_option(
170+
"-q",
171+
"--queue",
172+
action="append",
173+
dest="workflows_queue",
174+
default=[],
175+
metavar="[NAME=]QUEUE",
176+
help=(
177+
"Consume messages from QUEUE instead of the default queue NAME. "
178+
"NAME may be omitted for services that consume a single queue. "
179+
"May be specified multiple times. Defaults can be set in the "
180+
f"{QUEUE_ENVIRONMENT_VARIABLE} environment variable, as a comma or "
181+
"space separated list."
182+
),
183+
)
184+
140185
# Parse command line options
141186
(options, args) = parser.parse_args(cmdline_args)
142187

@@ -208,6 +253,25 @@ def on_transport_preparation_hook() -> CommonTransport:
208253
if add_metrics_option and options.metrics:
209254
kwargs["environment"]["metrics"] = {"port": options.metrics_port}
210255

256+
# Queue overrides from the command line take precedence over those
257+
# from the environment variable
258+
try:
259+
env_queue, env_queues = _parse_queue_overrides(
260+
os.environ.get(QUEUE_ENVIRONMENT_VARIABLE, "")
261+
.replace(",", " ")
262+
.split(),
263+
QUEUE_ENVIRONMENT_VARIABLE,
264+
)
265+
cli_queue, cli_queues = _parse_queue_overrides(
266+
getattr(options, "workflows_queue", []), "--queue"
267+
)
268+
except ValueError as e:
269+
parser.error(str(e))
270+
if queue := cli_queue or env_queue:
271+
kwargs["environment"]["queue"] = queue
272+
if queues := {**env_queues, **cli_queues}:
273+
kwargs["environment"]["queues"] = queues
274+
211275
if options.liveness:
212276
kwargs["environment"]["liveness"] = {
213277
"port": options.liveness_port,

‎src/workflows/services/common_service.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import workflows
2121
import workflows.logging
22+
import workflows.recipe
2223
from workflows.transport.common_transport import CommonTransport, MessageCallback
2324
from workflows.transport.middleware.otel_tracing import OTELTracingMiddleware
2425

@@ -168,6 +169,7 @@ def __init__(self, *, environment: dict[str, Any] | None = None):
168169
)
169170
self._idle_callback: Callable[[], None] | None = None
170171
self._idle_time: float | None = None
172+
self.__subscribed_channels: list[str] = []
171173
self.start_kwargs = {}
172174

173175
# Logger will be overwritten in start() function
@@ -406,6 +408,72 @@ def _register_idle(self, idle_time: float, callback: Callable[[], None]) -> None
406408
self._idle_callback = callback
407409
self._idle_time = idle_time
408410

411+
def _resolve_queue(self, channel: str) -> str:
412+
"""Return the name of the queue to subscribe to in place of channel.
413+
414+
Overrides are read from the ``environment`` dictionary passed to the
415+
service constructor (``self._environment``), not from operating system
416+
environment variables or the recipe environment. The service starter
417+
fills these entries from its --queue option and the WORKFLOWS_QUEUE
418+
environment variable:
419+
420+
- ``environment["queues"]``, a mapping of default queue names to the
421+
queue names to use instead
422+
- ``environment["queue"]``, a single queue name. This is only accepted
423+
for services subscribing to a single queue.
424+
"""
425+
if override := (self._environment.get("queues") or {}).get(channel):
426+
return override
427+
if override := self._environment.get("queue"):
428+
if self.__subscribed_channels:
429+
raise RuntimeError(
430+
f"Queue override {override!r} is ambiguous as this service "
431+
"subscribes to multiple queues. Use --queue NAME=QUEUE instead."
432+
)
433+
return override
434+
return channel
435+
436+
def subscribe(
437+
self,
438+
channel: str,
439+
callback: Callable[..., Any],
440+
*,
441+
acknowledgement: bool = True,
442+
**kwargs: Any,
443+
) -> int:
444+
"""Subscribe to a queue of recipe messages.
445+
446+
Messages are unpacked with :func:`workflows.recipe.wrap_subscribe`, so
447+
the callback receives ``(recipe_wrapper, header, message)``. Use
448+
``self.transport.subscribe`` for a subscription without recipe handling.
449+
450+
The queue name can be overridden when the service is started, with the
451+
--queue option or the WORKFLOWS_QUEUE environment variable, see
452+
:meth:`_resolve_queue`. By default messages must be acknowledged, and
453+
log messages are extended with the recipe ID.
454+
455+
Args:
456+
channel: Default name of the queue to subscribe to.
457+
callback: Function to be called when messages are received.
458+
acknowledgement: Whether received messages must be acknowledged.
459+
**kwargs: Passed on to :func:`workflows.recipe.wrap_subscribe`.
460+
461+
Returns:
462+
A unique subscription ID
463+
"""
464+
queue_name = self._resolve_queue(channel)
465+
if queue_name != channel:
466+
self.log.info("Subscribing to queue %s in place of %s", queue_name, channel)
467+
self.__subscribed_channels.append(channel)
468+
kwargs.setdefault("log_extender", self.extend_log)
469+
return workflows.recipe.wrap_subscribe(
470+
self.transport,
471+
queue_name,
472+
callback,
473+
acknowledgement=acknowledgement,
474+
**kwargs,
475+
)
476+
409477
def __update_service_status(self, statuscode: int) -> None:
410478
"""Set the internal status of the service object, and notify frontend."""
411479
if self.__service_status != statuscode:

‎tests/contrib/test_start_service.py‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ def test_script_initialises_transport_and_starts_frontend(
3030
mock_options.service = "someservice"
3131
mock_options.transport = mock.sentinel.transport
3232
mock_options.liveness = False
33+
mock_options.workflows_queue = []
3334
mock_parser.return_value.parse_args.return_value = (mock_options, mock.Mock())
3435
mock_services.get_known_services.return_value = {"SomeService": None}
3536

@@ -59,6 +60,7 @@ def test_add_metrics_option(mock_services, mock_frontend, mock_tlookup, mock_par
5960
mock_options.service = "someservice"
6061
mock_options.transport = mock.sentinel.transport
6162
mock_options.liveness = False
63+
mock_options.workflows_queue = []
6264
mock_options.metrics = False
6365
mock_options.metrics_port = 4242
6466
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
8991
transport=mock.ANY,
9092
environment={"metrics": {"port": 4242}},
9193
)
94+
95+
96+
@pytest.mark.parametrize(
97+
"args,environment",
98+
[
99+
([], {}),
100+
(["-q", "other"], {"queue": "other"}),
101+
(
102+
["--queue", "a=a.test", "-q", "b=b.test"],
103+
{"queues": {"a": "a.test", "b": "b.test"}},
104+
),
105+
(
106+
["-q", "a=a.test", "-q", "other"],
107+
{"queue": "other", "queues": {"a": "a.test"}},
108+
),
109+
],
110+
)
111+
@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
112+
@mock.patch("workflows.contrib.start_service.workflows.frontend")
113+
@mock.patch("workflows.contrib.start_service.workflows.services")
114+
def test_queue_option(mock_services, mock_frontend, mock_tlookup, args, environment):
115+
mock_services.get_known_services.return_value = {"SomeService": None}
116+
workflows.contrib.start_service.ServiceStarter().run(
117+
cmdline_args=["-s", "SomeService", *args]
118+
)
119+
assert mock_frontend.Frontend.call_args.kwargs["environment"] == environment
120+
121+
122+
@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
123+
@mock.patch("workflows.contrib.start_service.workflows.frontend")
124+
@mock.patch("workflows.contrib.start_service.workflows.services")
125+
def test_queue_option_allows_one_unnamed_queue(
126+
mock_services, mock_frontend, mock_tlookup, capsys
127+
):
128+
mock_services.get_known_services.return_value = {"SomeService": None}
129+
with pytest.raises(SystemExit):
130+
workflows.contrib.start_service.ServiceStarter().run(
131+
cmdline_args=["-s", "SomeService", "-q", "a", "-q", "b"]
132+
)
133+
assert "Only one --queue entry" in capsys.readouterr().err
134+
mock_frontend.Frontend.assert_not_called()
135+
136+
137+
@pytest.mark.parametrize(
138+
"variable,args,environment",
139+
[
140+
("other", [], {"queue": "other"}),
141+
("a=a.env, b=b.env", [], {"queues": {"a": "a.env", "b": "b.env"}}),
142+
(
143+
"a=a.env b=b.env other",
144+
[],
145+
{"queue": "other", "queues": {"a": "a.env", "b": "b.env"}},
146+
),
147+
# Command line options override the environment variable
148+
(
149+
"a=a.env,b=b.env",
150+
["-q", "a=a.cli"],
151+
{"queues": {"a": "a.cli", "b": "b.env"}},
152+
),
153+
("env", ["-q", "cli"], {"queue": "cli"}),
154+
("env", ["-q", "a=a.cli"], {"queue": "env", "queues": {"a": "a.cli"}}),
155+
("", [], {}),
156+
],
157+
)
158+
@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
159+
@mock.patch("workflows.contrib.start_service.workflows.frontend")
160+
@mock.patch("workflows.contrib.start_service.workflows.services")
161+
def test_queue_environment_variable(
162+
mock_services, mock_frontend, mock_tlookup, monkeypatch, variable, args, environment
163+
):
164+
monkeypatch.setenv("WORKFLOWS_QUEUE", variable)
165+
mock_services.get_known_services.return_value = {"SomeService": None}
166+
workflows.contrib.start_service.ServiceStarter().run(
167+
cmdline_args=["-s", "SomeService", *args]
168+
)
169+
assert mock_frontend.Frontend.call_args.kwargs["environment"] == environment
170+
171+
172+
@pytest.mark.parametrize(
173+
"variable,args,error",
174+
[
175+
("a b", [], "Only one WORKFLOWS_QUEUE entry"),
176+
("a=", [], "Missing queue name in WORKFLOWS_QUEUE entry 'a='"),
177+
("", ["-q", "a="], "Missing queue name in --queue entry 'a='"),
178+
],
179+
)
180+
@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
181+
@mock.patch("workflows.contrib.start_service.workflows.frontend")
182+
@mock.patch("workflows.contrib.start_service.workflows.services")
183+
def test_queue_overrides_rejected(
184+
mock_services,
185+
mock_frontend,
186+
mock_tlookup,
187+
monkeypatch,
188+
capsys,
189+
variable,
190+
args,
191+
error,
192+
):
193+
monkeypatch.setenv("WORKFLOWS_QUEUE", variable)
194+
mock_services.get_known_services.return_value = {"SomeService": None}
195+
with pytest.raises(SystemExit):
196+
workflows.contrib.start_service.ServiceStarter().run(
197+
cmdline_args=["-s", "SomeService", *args]
198+
)
199+
assert error in capsys.readouterr().err
200+
mock_frontend.Frontend.assert_not_called()
201+
202+
203+
class StarterWithOwnQueueOption(workflows.contrib.start_service.ServiceStarter):
204+
"""A subclass defining its own queue option, as dlstbx did before workflows"""
205+
206+
def __init__(self, *option_strings: str):
207+
self.option_strings = option_strings or ("-q", "--queue")
208+
209+
def on_parser_preparation(self, parser):
210+
parser.add_option(*self.option_strings, dest="queue", default=None)
211+
212+
def on_parsing(self, options, args):
213+
self.options = options
214+
215+
def before_frontend_construction(self, kwargs):
216+
if self.options.queue:
217+
kwargs["environment"]["queue"] = self.options.queue
218+
return kwargs
219+
220+
221+
@pytest.mark.parametrize("option_strings", [(), ("--queue",), ("-q", "--quiet")])
222+
@mock.patch("workflows.contrib.start_service.workflows.transport.lookup")
223+
@mock.patch("workflows.contrib.start_service.workflows.frontend")
224+
@mock.patch("workflows.contrib.start_service.workflows.services")
225+
def test_queue_option_skipped_if_already_defined(
226+
mock_services, mock_frontend, mock_tlookup, monkeypatch, option_strings
227+
):
228+
monkeypatch.setenv("WORKFLOWS_QUEUE", "a=a.env")
229+
mock_services.get_known_services.return_value = {"SomeService": None}
230+
starter = StarterWithOwnQueueOption(*option_strings)
231+
flag = (option_strings or ("-q",))[0]
232+
starter.run(cmdline_args=["-s", "SomeService", flag, "own"])
233+
234+
# The subclass option is parsed as before, as a single value
235+
assert starter.options.queue == "own"
236+
assert not hasattr(starter.options, "workflows_queue")
237+
# The environment variable still applies
238+
assert mock_frontend.Frontend.call_args.kwargs["environment"] == {
239+
"queue": "own",
240+
"queues": {"a": "a.env"},
241+
}

‎tests/services/test_common_service.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -626,3 +626,48 @@ def test_transport_callbacks_are_processed_from_main_queue():
626626
service.SERVICE_STATUS_SHUTDOWN,
627627
service.SERVICE_STATUS_END,
628628
]
629+
630+
631+
@pytest.fixture
632+
def wrap_subscribe(mocker):
633+
return mocker.patch("workflows.recipe.wrap_subscribe")
634+
635+
636+
def test_subscribe_uses_default_queue(wrap_subscribe):
637+
service = CommonService()
638+
service.transport = mock.Mock()
639+
callback = mock.Mock()
640+
service.subscribe("some_queue", callback, exclusive=True)
641+
wrap_subscribe.assert_called_once_with(
642+
service.transport,
643+
"some_queue",
644+
callback,
645+
acknowledgement=True,
646+
log_extender=service.extend_log,
647+
exclusive=True,
648+
)
649+
650+
651+
def test_subscribe_with_named_overrides(wrap_subscribe):
652+
service = CommonService(environment={"queues": {"a": "a.test"}})
653+
service.transport = mock.Mock()
654+
service.subscribe("a", mock.Mock())
655+
service.subscribe("b", mock.Mock(), acknowledgement=False)
656+
assert [c.args[1] for c in wrap_subscribe.call_args_list] == ["a.test", "b"]
657+
assert wrap_subscribe.call_args.kwargs["acknowledgement"] is False
658+
659+
660+
def test_subscribe_bare_override_applies_to_single_queue_only(wrap_subscribe):
661+
service = CommonService(environment={"queue": "x.test"})
662+
service.transport = mock.Mock()
663+
service.subscribe("a", mock.Mock())
664+
assert wrap_subscribe.call_args.args[1] == "x.test"
665+
with pytest.raises(RuntimeError, match="ambiguous"):
666+
service.subscribe("b", mock.Mock())
667+
668+
669+
def test_subscribe_named_override_takes_precedence(wrap_subscribe):
670+
service = CommonService(environment={"queue": "x.test", "queues": {"a": "a.test"}})
671+
service.transport = mock.Mock()
672+
service.subscribe("a", mock.Mock())
673+
assert wrap_subscribe.call_args.args[1] == "a.test"

0 commit comments

Comments
 (0)