durable-worker 0.2.0__tar.gz → 0.3.0__tar.gz

{durable_worker-0.2.0 → durable_worker-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durable-worker
-Version: 0.2.0
+Version: 0.3.0
 Summary: Python worker SDK for nestjs-durable — run durable workflow steps in Python.
 Project-URL: Homepage, https://github.com/DavideCarvalho/nestjs-durable
 Project-URL: Repository, https://github.com/DavideCarvalho/nestjs-durable

{durable_worker-0.2.0 → durable_worker-0.3.0}/durable_worker/__init__.py

@@ -5,7 +5,16 @@ The wire protocol (task in, result out) is plain JSON and identical across langu
 same step name implemented here is callable from a TypeScript workflow via ``ctx.call``.
 """
 
+from .cancellation import Cancelled, CancellationRegistry
+from .routing import reply_target
 from .worker import FatalError, StepContext, Worker
 
-__all__ = ["Worker", "FatalError", "StepContext"]
-__version__ = "0.2.0"
+__all__ = [
+    "Worker",
+    "FatalError",
+    "StepContext",
+    "Cancelled",
+    "CancellationRegistry",
+    "reply_target",
+]
+__version__ = "0.3.0"

durable_worker-0.3.0/durable_worker/cancellation.py

@@ -0,0 +1,49 @@
+"""Cooperative cancellation for the Python worker.
+
+The orchestrator broadcasts a ``{"kind": "cancel", "runId": ...}`` control message across every
+instance (Redis pub/sub, same channel the TypeScript control plane uses). A long-running handler
+can't be force-killed, so cancellation is *cooperative*: a runner subscribes to that channel and
+records the cancelled run ids here, and the handler checks ``ctx.cancelled`` (or calls
+``ctx.raise_if_cancelled()``) at safe points to bail out early instead of finishing work whose
+result will be discarded.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Dict
+
+
+class Cancelled(Exception):
+    """Raised by ``StepContext.raise_if_cancelled`` when the run has been cancelled."""
+
+    def __init__(self, run_id: str) -> None:
+        super().__init__(f"run {run_id} was cancelled")
+        self.run_id = run_id
+
+
+class CancellationRegistry:
+    """Thread-safe set of cancelled run ids, fed by a control-channel subscription."""
+
+    def __init__(self) -> None:
+        self._cancelled: set[str] = set()
+        self._lock = threading.Lock()
+
+    def cancel(self, run_id: str) -> None:
+        with self._lock:
+            self._cancelled.add(run_id)
+
+    def is_cancelled(self, run_id: str) -> bool:
+        with self._lock:
+            return run_id in self._cancelled
+
+    def clear(self, run_id: str) -> None:
+        with self._lock:
+            self._cancelled.discard(run_id)
+
+    def on_control_message(self, msg: Dict[str, Any]) -> None:
+        """Feed a decoded control message; records the run id when it's a cancel."""
+        if isinstance(msg, dict) and msg.get("kind") == "cancel":
+            run_id = msg.get("runId")
+            if isinstance(run_id, str):
+                self.cancel(run_id)

durable_worker-0.3.0/durable_worker/redis_runner.py

@@ -0,0 +1,159 @@
+"""Run a :class:`Worker` against the BullMQ/Redis transport.
+
+Consumes the orchestrator's per-group tasks queue and publishes results on the shared results
+queue — the same queues a TypeScript ``BullMQTransport`` uses, so steps interoperate across
+languages. Requires the optional ``bullmq`` extra: ``pip install durable-worker[redis]``.
+
+It also subscribes to the orchestrator's control channel (``<prefix>-control``) so a long handler
+can observe cooperative cancellation via ``ctx.cancelled`` — the cross-language half of
+``engine.cancel``.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import socket
+from typing import Any, Callable, Dict, Optional
+
+from .cancellation import CancellationRegistry
+from .worker import Worker
+
+# Stable-ish id for the `from` field of control messages this worker publishes. It only has to
+# DIFFER from the engine instanceIds (so a dashboard engine doesn't treat our progress events as its
+# own echo and drop them) — host + pid is plenty and avoids importing a uuid/random dependency.
+_INSTANCE_ID = f"py-{socket.gethostname()}-{os.getpid()}"
+
+
+def _names(prefix: str, group: str) -> tuple[str, str]:
+    # Must match the TS BullMQTransport: '<prefix>-tasks-<group>' and '<prefix>-results'.
+    return f"{prefix}-tasks-{group}", f"{prefix}-results"
+
+
+def _control_channel(prefix: str) -> str:
+    # Mirrors BullMQTransport.controlChannel(): '<prefix>-control'.
+    return f"{prefix}-control"
+
+
+def _progress_message(task: Dict[str, Any], event: Dict[str, Any]) -> str:
+    """A control-plane `{kind:'event'}` carrying a `step.progress` EngineEvent — the same envelope a
+    TS engine publishes, so a dashboard engine re-delivers it to its subscribers (live-tail). Carries
+    the single just-emitted step event; the `at` mirrors EngineEvent's (ms; the TS side `new Date()`s it)."""
+    return json.dumps(
+        {
+            "kind": "event",
+            "from": _INSTANCE_ID,
+            "event": {
+                "type": "step.progress",
+                "runId": task.get("runId"),
+                "seq": task.get("seq"),
+                "name": task.get("name"),
+                "event": event,
+                "at": event.get("at"),
+            },
+        }
+    )
+
+
+async def run_redis_worker(
+    worker: Worker,
+    *,
+    group: str,
+    connection: str = "redis://localhost:6379",
+    prefix: str = "durable",
+    cancellation: Optional[CancellationRegistry] = None,
+) -> Any:
+    """Start a BullMQ worker that runs ``worker``'s handlers. Returns the bullmq Worker.
+
+    The returned worker runs in the background; ``await worker.close()`` to stop it. When a
+    :class:`CancellationRegistry` is given (one is created otherwise), the runner subscribes to the
+    control channel and feeds it, so handlers see ``ctx.cancelled``.
+    """
+
+    from bullmq import Queue as BullQueue  # imported lazily so the SDK works without bullmq
+    from bullmq import Worker as BullWorker
+
+    tasks_name, results_name = _names(prefix, group)
+    results = BullQueue(results_name, {"connection": connection})
+    registry = cancellation or CancellationRegistry()
+    await _subscribe_control(connection, prefix, registry)
+    publish_progress = await _progress_publisher(connection, prefix)
+
+    async def process(job: Any, _token: str) -> None:
+        on_event = _make_on_event(job.data, publish_progress) if publish_progress else None
+        result = await worker.aprocess_task(
+            job.data, is_cancelled=registry.is_cancelled, on_event=on_event
+        )
+        await results.add("result", result, {"removeOnComplete": True, "removeOnFail": True})
+
+    return BullWorker(tasks_name, process, {"connection": connection})
+
+
+async def _progress_publisher(
+    connection: str, prefix: str
+) -> Optional[Callable[[str], None]]:
+    """Build a thread-safe `publish(message)` that PUBLISHes on the control channel from the running
+    loop. Returns None if redis pub/sub isn't available (then `step.progress` streaming is simply off
+    — the events still ride back on the final result). The returned callable is safe to call from a
+    handler's executor thread: it hops back onto the loop via ``call_soon_threadsafe`` (the aioredis
+    client is bound to this loop, so publishing must happen there, not from the worker thread)."""
+    try:
+        import redis.asyncio as aioredis  # lazy: only needed when streaming progress
+    except ImportError:
+        return None
+
+    client = aioredis.from_url(connection)
+    channel = _control_channel(prefix)
+    loop = asyncio.get_running_loop()
+
+    def publish(message: str) -> None:
+        async def _send() -> None:
+            try:
+                await client.publish(channel, message)
+            except Exception:  # noqa: BLE001 — live-tail is best-effort; never fail the step
+                pass
+
+        try:
+            loop.call_soon_threadsafe(lambda: loop.create_task(_send()))
+        except RuntimeError:
+            pass  # loop is closing/closed — drop the live event (the result still carries it)
+
+    return publish
+
+
+def _make_on_event(
+    task: Dict[str, Any], publish: Callable[[str], None]
+) -> Callable[[Dict[str, Any]], None]:
+    """Per-task sink: turn each step event into a `step.progress` control message and publish it."""
+
+    def on_event(event: Dict[str, Any]) -> None:
+        publish(_progress_message(task, event))
+
+    return on_event
+
+
+async def _subscribe_control(
+    connection: str, prefix: str, registry: CancellationRegistry
+) -> None:
+    """Best-effort: subscribe to the control channel and feed cancellations into ``registry``.
+    No-op (logged) if redis pub/sub isn't available — cancellation just won't be observed."""
+    try:
+        import redis.asyncio as aioredis  # lazy: only needed for cooperative cancellation
+    except ImportError:
+        return
+
+    client = aioredis.from_url(connection)
+    pubsub = client.pubsub()
+    await pubsub.subscribe(_control_channel(prefix))
+
+    async def listen() -> None:
+        async for message in pubsub.listen():
+            if message.get("type") != "message":
+                continue
+            try:
+                registry.on_control_message(json.loads(message["data"]))
+            except (ValueError, TypeError):
+                pass  # ignore malformed control messages
+
+    asyncio.create_task(listen())

durable_worker-0.3.0/durable_worker/routing.py

@@ -0,0 +1,27 @@
+"""Reply routing for a worker that consumes several transports (multi-broker failover).
+
+The orchestrator stamps each task with the id of the transport it was dispatched on
+(``task["transport"]``). A worker that consumes more than one broker must send the *result* back on
+the same broker — so failover is symmetric and the worker never picks a transport itself, it just
+replies where it was told. This is a pure lookup over a ``{id: reply_target}`` mapping with a sane
+fallback for the single-transport / unknown-id case.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, TypeVar
+
+T = TypeVar("T")
+
+
+def reply_target(task: Dict[str, Any], targets: Dict[str, T]) -> Optional[T]:
+    """The reply target for ``task`` by its ``transport`` id; falls back to the lone target when the
+    id is absent (single transport) or unknown. Returns ``None`` only if ``targets`` is empty."""
+    if not targets:
+        return None
+    transport_id = task.get("transport")
+    if isinstance(transport_id, str) and transport_id in targets:
+        return targets[transport_id]
+    if len(targets) == 1:
+        return next(iter(targets.values()))
+    return None

{durable_worker-0.2.0 → durable_worker-0.3.0}/durable_worker/worker.py

@@ -43,8 +43,47 @@ class StepContext:
             return {"context": {...}}
     """
 
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        run_id: Optional[str] = None,
+        is_cancelled: Optional[Callable[[str], bool]] = None,
+        seq: Optional[int] = None,
+        on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
+    ) -> None:
         self.events: List[Dict[str, Any]] = []
+        self._run_id = run_id
+        self._is_cancelled = is_cancelled
+        # Step seq + a sink for live progress: when a transport supplies ``on_event``, every event is
+        # ALSO handed to it as it happens (e.g. published on the control plane as a ``step.progress``),
+        # so a dashboard tails a long step line-by-line instead of waiting for the final result.
+        self._seq = seq
+        self._on_event = on_event
+        # The sub-process a handler is currently inside (set via ``process``), stamped onto log lines
+        # so the dashboard groups a fan-out step's trail per sub-process.
+        self._process: Optional[str] = None
+
+    def process(self, name: Optional[str]) -> None:
+        """Mark the sub-process whose work is now running; subsequent log lines are tagged with it
+        (until changed or cleared with ``None``). A no-op for ``sub()`` outcome rows, which name
+        themselves. Lets a fan-out handler attribute its log trail to each sub-process."""
+        self._process = name
+
+    @property
+    def cancelled(self) -> bool:
+        """True once the run has been cancelled (a runner subscribed to the control channel marks
+        it). Check this at safe points in a long handler and return/raise to bail out early."""
+        return bool(
+            self._is_cancelled is not None
+            and self._run_id is not None
+            and self._is_cancelled(self._run_id)
+        )
+
+    def raise_if_cancelled(self) -> None:
+        """Raise :class:`Cancelled` if the run has been cancelled — a one-liner abort for handlers."""
+        if self.cancelled:
+            from .cancellation import Cancelled
+
+            raise Cancelled(self._run_id or "")
 
     def _emit(
         self,
@@ -59,9 +98,19 @@ class StepContext:
             event["name"] = name
         if status is not None:
             event["status"] = status
+        # Tag a log line (no outcome status) with the sub-process it was emitted inside, if any.
+        elif self._process is not None:
+            event["process"] = self._process
         if data is not None:
             event["data"] = data
         self.events.append(event)
+        # Live progress: hand the event to the transport's sink so it streams now (best-effort — a
+        # broken sink must never fail the handler; the event is already captured for the result).
+        if self._on_event is not None:
+            try:
+                self._on_event(event)
+            except Exception:  # noqa: BLE001 — live-tail is best-effort observability
+                pass
 
     def debug(self, message: str, data: Any = None) -> None:
         self._emit("debug", message, data=data)
@@ -112,18 +161,25 @@ class Worker:
     def handles(self, name: str) -> bool:
         return name in self._handlers
 
-    def process_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
+    def process_task(
+        self,
+        task: Dict[str, Any],
+        is_cancelled: Optional[Callable[[str], bool]] = None,
+        on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
+    ) -> Dict[str, Any]:
         """Run the handler for ``task`` and return a wire-format result.
 
         Pure and synchronous from the caller's view (async handlers are awaited internally), so a
-        transport can simply ``result = worker.process_task(task); send(result)``.
+        transport can simply ``result = worker.process_task(task); send(result)``. Pass
+        ``is_cancelled`` so the handler's ``ctx.cancelled`` reflects cooperative cancellation, and
+        ``on_event`` to receive each step event live (for ``step.progress`` streaming).
         """
 
         base = self._base(task)
         handler = self._handlers.get(task["name"])
         if handler is None:
             return self._no_handler(base, task["name"])
-        ctx = StepContext()
+        ctx = self._ctx(task, is_cancelled, on_event)
         try:
             output = self._invoke(handler, task.get("input"), ctx)
             if inspect.isawaitable(output):
@@ -132,7 +188,12 @@ class Worker:
         except Exception as err:  # noqa: BLE001
             return self._failure(base, err, ctx)
 
-    async def aprocess_task(self, task: Dict[str, Any]) -> Dict[str, Any]:
+    async def aprocess_task(
+        self,
+        task: Dict[str, Any],
+        is_cancelled: Optional[Callable[[str], bool]] = None,
+        on_event: Optional[Callable[[Dict[str, Any]], None]] = None,
+    ) -> Dict[str, Any]:
         """Async variant — awaits async handlers in the current loop. Use from a transport that
         already runs inside an event loop (e.g. the BullMQ runner)."""
 
@@ -140,7 +201,7 @@ class Worker:
         handler = self._handlers.get(task["name"])
         if handler is None:
             return self._no_handler(base, task["name"])
-        ctx = StepContext()
+        ctx = self._ctx(task, is_cancelled, on_event)
         try:
             output = self._invoke(handler, task.get("input"), ctx)
             if inspect.isawaitable(output):
@@ -149,6 +210,19 @@ class Worker:
         except Exception as err:  # noqa: BLE001
             return self._failure(base, err, ctx)
 
+    @staticmethod
+    def _ctx(
+        task: Dict[str, Any],
+        is_cancelled: Optional[Callable[[str], bool]],
+        on_event: Optional[Callable[[Dict[str, Any]], None]],
+    ) -> StepContext:
+        return StepContext(
+            run_id=task.get("runId"),
+            is_cancelled=is_cancelled,
+            seq=task.get("seq"),
+            on_event=on_event,
+        )
+
     @staticmethod
     def _invoke(handler: Handler, input_: Any, ctx: StepContext) -> Any:
         """Call ``handler`` with the step input, passing ``ctx`` only if it declares a second

{durable_worker-0.2.0 → durable_worker-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "durable-worker"
-version = "0.2.0"
+version = "0.3.0"
 description = "Python worker SDK for nestjs-durable — run durable workflow steps in Python."
 readme = "README.md"
 requires-python = ">=3.9"

durable_worker-0.3.0/tests/test_cancellation.py

@@ -0,0 +1,68 @@
+import unittest
+
+from durable_worker import CancellationRegistry, Cancelled, Worker
+
+
+def task(name="long.job", **over):
+    base = {"runId": "r1", "seq": 0, "name": name, "stepId": "r1:0", "input": {}, "attempt": 1}
+    base.update(over)
+    return base
+
+
+class CancellationRegistryTest(unittest.TestCase):
+    def test_records_a_cancel_control_message(self):
+        reg = CancellationRegistry()
+        self.assertFalse(reg.is_cancelled("r1"))
+        reg.on_control_message({"kind": "cancel", "runId": "r1", "from": "x"})
+        self.assertTrue(reg.is_cancelled("r1"))
+
+    def test_ignores_non_cancel_messages(self):
+        reg = CancellationRegistry()
+        reg.on_control_message({"kind": "event", "event": {}})
+        reg.on_control_message({"nonsense": True})
+        self.assertFalse(reg.is_cancelled("r1"))
+
+
+class StepContextCancellationTest(unittest.TestCase):
+    def test_handler_sees_cancellation_via_ctx(self):
+        reg = CancellationRegistry()
+        worker = Worker()
+        observed = {}
+
+        @worker.step("long.job")
+        def run(_data, ctx):
+            observed["before"] = ctx.cancelled
+            reg.cancel("r1")
+            observed["after"] = ctx.cancelled
+            return "done"
+
+        worker.process_task(task(), is_cancelled=reg.is_cancelled)
+        self.assertEqual(observed, {"before": False, "after": True})
+
+    def test_raise_if_cancelled_aborts_the_step(self):
+        reg = CancellationRegistry()
+        reg.cancel("r1")
+        worker = Worker()
+
+        @worker.step("long.job")
+        def run(_data, ctx):
+            ctx.raise_if_cancelled()
+            return "should not get here"
+
+        result = worker.process_task(task(), is_cancelled=reg.is_cancelled)
+        self.assertEqual(result["status"], "failed")
+        self.assertIn("cancelled", result["error"]["message"])
+
+    def test_cancelled_is_false_without_a_source(self):
+        worker = Worker()
+
+        @worker.step("long.job")
+        def run(_data, ctx):
+            return ctx.cancelled
+
+        result = worker.process_task(task())
+        self.assertEqual(result["output"], False)
+
+
+if __name__ == "__main__":
+    unittest.main()

durable_worker-0.3.0/tests/test_routing.py

@@ -0,0 +1,23 @@
+import unittest
+
+from durable_worker import reply_target
+
+
+class ReplyTargetTest(unittest.TestCase):
+    def test_routes_by_task_transport_id(self):
+        targets = {"redis": "redis-results", "sqs": "sqs-results"}
+        self.assertEqual(reply_target({"transport": "sqs"}, targets), "sqs-results")
+        self.assertEqual(reply_target({"transport": "redis"}, targets), "redis-results")
+
+    def test_falls_back_to_the_lone_target_when_id_is_absent(self):
+        self.assertEqual(reply_target({}, {"only": "x"}), "x")
+
+    def test_unknown_id_with_multiple_targets_is_unroutable(self):
+        self.assertIsNone(reply_target({"transport": "nats"}, {"redis": "a", "sqs": "b"}))
+
+    def test_empty_targets_is_none(self):
+        self.assertIsNone(reply_target({"transport": "redis"}, {}))
+
+
+if __name__ == "__main__":
+    unittest.main()

{durable_worker-0.2.0 → durable_worker-0.3.0}/tests/test_worker.py

@@ -136,6 +136,48 @@ class WorkerTest(unittest.TestCase):
         result = worker.process_task(task())
         self.assertEqual(result["events"][0]["message"], "async")
 
+    def test_logs_are_tagged_with_the_current_process(self):
+        worker = Worker()
+
+        @worker.step("payments.charge-card")
+        def charge(_data, ctx: StepContext):
+            ctx.process("proc-a")
+            ctx.info("inside a")
+            ctx.sub("proc-a", "ok")
+            ctx.process("proc-b")
+            ctx.debug("inside b")
+            ctx.process(None)
+            ctx.info("step-level")
+            return {"ok": True}
+
+        # process() itself emits nothing, so the events are: [info a, sub a, debug b, info step-level]
+        events = worker.process_task(task())["events"]
+        # log lines carry the sub-process that was running when they fired...
+        self.assertEqual(events[0]["process"], "proc-a")
+        self.assertEqual(events[2]["process"], "proc-b")
+        # ...the sub-process OUTCOME row names itself (no `process` tag)...
+        self.assertNotIn("process", events[1])
+        self.assertEqual(events[1]["status"], "ok")
+        # ...and a log after process(None) is untagged (step-level).
+        self.assertNotIn("process", events[3])
+
+    def test_on_event_streams_each_event_as_it_is_emitted(self):
+        worker = Worker()
+        live = []
+
+        @worker.step("payments.charge-card")
+        def charge(_data, ctx: StepContext):
+            ctx.info("first")
+            # the live sink saw the event the instant it was emitted, before the step returned
+            self.assertEqual([e["message"] for e in live], ["first"])
+            ctx.sub("proc-a", "ok")
+            return {"ok": True}
+
+        result = worker.process_task(task(), on_event=live.append)
+        self.assertEqual([e["message"] for e in live], ["first", "proc-a"])
+        # the same events still ride back on the result (live-tail doesn't replace the final record)
+        self.assertEqual(len(result["events"]), 2)
+
 
 if __name__ == "__main__":
     unittest.main()

durable_worker-0.2.0/durable_worker/redis_runner.py

@@ -1,42 +0,0 @@
-"""Run a :class:`Worker` against the BullMQ/Redis transport.
-
-Consumes the orchestrator's per-group tasks queue and publishes results on the shared results
-queue — the same queues a TypeScript ``BullMQTransport`` uses, so steps interoperate across
-languages. Requires the optional ``bullmq`` extra: ``pip install durable-worker[redis]``.
-"""
-
-from __future__ import annotations
-
-from typing import Any
-
-from .worker import Worker
-
-
-def _names(prefix: str, group: str) -> tuple[str, str]:
-    # Must match the TS BullMQTransport: '<prefix>-tasks-<group>' and '<prefix>-results'.
-    return f"{prefix}-tasks-{group}", f"{prefix}-results"
-
-
-async def run_redis_worker(
-    worker: Worker,
-    *,
-    group: str,
-    connection: str = "redis://localhost:6379",
-    prefix: str = "durable",
-) -> Any:
-    """Start a BullMQ worker that runs ``worker``'s handlers. Returns the bullmq Worker.
-
-    The returned worker runs in the background; ``await worker.close()`` to stop it.
-    """
-
-    from bullmq import Queue as BullQueue  # imported lazily so the SDK works without bullmq
-    from bullmq import Worker as BullWorker
-
-    tasks_name, results_name = _names(prefix, group)
-    results = BullQueue(results_name, {"connection": connection})
-
-    async def process(job: Any, _token: str) -> None:
-        result = await worker.aprocess_task(job.data)
-        await results.add("result", result, {"removeOnComplete": True, "removeOnFail": True})
-
-    return BullWorker(tasks_name, process, {"connection": connection})