bunqueue-client 0.1.5__py3-none-any.whl

bunqueue/flow.py

@@ -0,0 +1,257 @@
+"""FlowProducer: parent/child job trees, chains, and fan-in flows.
+
+Mirrors the TS FlowProducer (TCP mode): children are pushed BEFORE their
+parent; the parent carries dependsOn/childrenIds; each child is then linked
+via UpdateParent. On any failure every already-created job is cancelled
+(atomic rollback).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Sequence
+
+from .connection import Connection, TlsOption, _compact
+from .errors import CommandError
+from .job import Job
+from .options import job_options, job_payload
+from .telemetry import TelemetryHandler
+
+
+class FlowNode:
+    """Result node: the created job plus its (optional) children nodes."""
+
+    def __init__(self, job: Job, children: Optional[List["FlowNode"]] = None) -> None:
+        self.job = job
+        self.children = children or []
+
+
+class FlowProducer:
+    """Create dependent job hierarchies across one or more queues.
+
+    Flow tree entries: ``{"name", "queueName", "data", "opts", "children": [...]}``
+    (same shape as BullMQ / the TS client).
+    """
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 6789,
+        token: Optional[str] = None,
+        tls: TlsOption = None,
+        connection: Optional[Connection] = None,
+        on_telemetry: Optional[TelemetryHandler] = None,
+    ) -> None:
+        self.connection = connection or Connection(
+            host=host, port=port, token=token, tls=tls, on_telemetry=on_telemetry
+        )
+        self._owns_connection = connection is None
+
+    # ------------------------------------------------------------------- api
+
+    def add(self, flow: Dict[str, Any], opts: Optional[Dict[str, Any]] = None) -> FlowNode:
+        """Add a flow tree. Children are processed BEFORE their parent."""
+        created: List[str] = []
+        try:
+            return self._add_node(flow, None, created, opts or {})
+        except Exception:
+            self._cleanup(created)
+            raise
+
+    def add_bulk(self, flows: Sequence[Dict[str, Any]]) -> List[FlowNode]:
+        results: List[FlowNode] = []
+        created: List[str] = []
+        try:
+            for flow in flows:
+                results.append(self._add_node(flow, None, created, {}))
+            return results
+        except Exception:
+            self._cleanup(created)
+            raise
+
+    def add_chain(self, steps: Sequence[Dict[str, Any]]) -> List[str]:
+        """Sequential chain: step[0] -> step[1] -> ... (each depends on the previous)."""
+        job_ids: List[str] = []
+        prev_id: Optional[str] = None
+        try:
+            for step in steps:
+                data = {"__flowParentId": prev_id, **(step.get("data") or {})}
+                job_id = self._push(
+                    step["queueName"],
+                    job_payload(step["name"], data),
+                    step.get("opts") or {},
+                    depends_on=[prev_id] if prev_id else None,
+                )
+                job_ids.append(job_id)
+                prev_id = job_id
+            return job_ids
+        except Exception:
+            self._cleanup(job_ids)
+            raise
+
+    def add_bulk_then(
+        self, parallel: Sequence[Dict[str, Any]], final: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Fan-in: N independent jobs converging into one final job.
+
+        The final job records the parallel ids as childrenIds so its processor
+        can read their results via ``get_children_values``.
+        """
+        parallel_ids: List[str] = []
+        try:
+            for step in parallel:
+                parallel_ids.append(
+                    self._push(
+                        step["queueName"],
+                        job_payload(step["name"], step.get("data")),
+                        step.get("opts") or {},
+                    )
+                )
+            final_data = {"__flowParentIds": parallel_ids, **(final.get("data") or {})}
+            final_id = self._push_with_children(
+                final["queueName"],
+                job_payload(final["name"], final_data),
+                final.get("opts") or {},
+                child_ids=parallel_ids,
+                parent_id=None,
+            )
+            return {"parallel_ids": parallel_ids, "final_id": final_id}
+        except Exception:
+            self._cleanup(parallel_ids)
+            raise
+
+    def get_flow(self, job_id: str, depth: Optional[int] = None) -> Optional[FlowNode]:
+        """Reconstruct a flow tree from a root job id (GetJob recursion).
+
+        ``depth`` limits recursion; ``None`` means unlimited (parity with the
+        TS client's ``Number.POSITIVE_INFINITY`` default). A missing job — the
+        root or any child removed via removeOnComplete/cancel — yields ``None``
+        and is skipped, returning the surviving partial tree instead of raising.
+        A ``visited`` set guards against cycles in ``childrenIds`` so unlimited
+        depth cannot recurse forever.
+        """
+        return self._get_flow(job_id, depth, set())
+
+    def _get_flow(
+        self, job_id: str, depth: Optional[int], visited: set
+    ) -> Optional[FlowNode]:
+        if job_id in visited:
+            return None  # cycle guard: childrenIds already on the current path
+        visited.add(job_id)
+        try:
+            response = self.connection.call({"cmd": "GetJob", "id": job_id})
+        except CommandError as exc:
+            if "not found" in str(exc).lower():
+                return None  # root or a since-removed child -> skip
+            raise  # a real server error must not masquerade as a missing child
+        raw = response.get("job")
+        if not raw:
+            return None
+        job = Job(raw, self.connection)
+        children: List[FlowNode] = []
+        if depth is None or depth > 0:
+            next_depth = None if depth is None else depth - 1
+            for child_id in job.children_ids:
+                child = self._get_flow(child_id, next_depth, visited)
+                if child:
+                    children.append(child)
+        return FlowNode(job, children)
+
+    def close(self) -> None:
+        if self._owns_connection:
+            self.connection.close()
+
+    def __enter__(self) -> "FlowProducer":
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.close()
+
+    # ------------------------------------------------------------- internals
+
+    def _add_node(
+        self,
+        node: Dict[str, Any],
+        parent_ref: Optional[Dict[str, str]],
+        created: List[str],
+        flow_opts: Dict[str, Any],
+    ) -> FlowNode:
+        child_nodes: List[FlowNode] = []
+        child_ids: List[str] = []
+
+        for child in node.get("children") or []:
+            child_node = self._add_node(
+                child, {"id": "pending", "queue": node["queueName"]}, created, flow_opts
+            )
+            child_nodes.append(child_node)
+            child_ids.append(child_node.job.id)
+
+        queue_defaults = (flow_opts.get("queues_options") or {}).get(node["queueName"]) or {}
+        merged_opts = {**queue_defaults, **(node.get("opts") or {})}
+
+        data: Dict[str, Any] = dict(node.get("data") or {})
+        if parent_ref:
+            data["__parentId"] = parent_ref["id"]
+            data["__parentQueue"] = parent_ref["queue"]
+        if child_ids:
+            data["__childrenIds"] = child_ids
+
+        job_id = self._push_with_children(
+            node["queueName"],
+            job_payload(node["name"], data),
+            merged_opts,
+            child_ids=child_ids,
+            parent_id=parent_ref["id"] if parent_ref else None,
+        )
+        created.append(job_id)
+
+        job = Job({"id": job_id, "queue": node["queueName"], "data": data}, self.connection)
+        return FlowNode(job, child_nodes)
+
+    def _push(
+        self,
+        queue: str,
+        data: Dict[str, Any],
+        opts: Dict[str, Any],
+        depends_on: Optional[List[str]] = None,
+    ) -> str:
+        command = {
+            "cmd": "PUSH",
+            "queue": queue,
+            "data": data,
+            **job_options(**opts),
+        }
+        if depends_on:
+            command["dependsOn"] = depends_on
+        response = self.connection.call(_compact(command))
+        return str(response.get("id"))
+
+    def _push_with_children(
+        self,
+        queue: str,
+        data: Dict[str, Any],
+        opts: Dict[str, Any],
+        child_ids: List[str],
+        parent_id: Optional[str],
+    ) -> str:
+        command = {
+            "cmd": "PUSH",
+            "queue": queue,
+            "data": data,
+            **job_options(**opts),
+            "parentId": parent_id,
+            "childrenIds": child_ids or None,
+            "dependsOn": child_ids or None,
+        }
+        response = self.connection.call(_compact(command))
+        job_id = str(response.get("id"))
+        # Link children to their real parent id (placeholder was 'pending')
+        for child_id in child_ids:
+            self.connection.call({"cmd": "UpdateParent", "childId": child_id, "parentId": job_id})
+        return job_id
+
+    def _cleanup(self, job_ids: List[str]) -> None:
+        for job_id in job_ids:
+            try:
+                self.connection.call({"cmd": "Cancel", "id": job_id})
+            except Exception:  # noqa: BLE001 - best-effort rollback
+                pass

bunqueue/job.py

@@ -0,0 +1,179 @@
+"""Job wrapper: read view over the server job record + per-job operations.
+
+Parity with the TS client's Job object (TCP mode): progress, logs, lock
+management, state transitions, children values.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Optional
+
+from .connection import Connection, _compact
+
+
+class Job:
+    """A job as seen by producers (from queries) or workers (with lock token)."""
+
+    def __init__(
+        self,
+        raw: Dict[str, Any],
+        connection: Optional[Connection] = None,
+        token: Optional[str] = None,
+        on_progress: Optional[Callable[["Job", Any], None]] = None,
+    ) -> None:
+        self.raw = raw
+        self.token = token
+        self._conn = connection
+        self._on_progress = on_progress
+
+    # ------------------------------------------------------------ properties
+
+    @property
+    def id(self) -> str:
+        return str(self.raw.get("id"))
+
+    @property
+    def queue(self) -> str:
+        return str(self.raw.get("queue"))
+
+    @property
+    def data(self) -> Any:
+        return self.raw.get("data")
+
+    @property
+    def name(self) -> Optional[str]:
+        """Job name (stored inside ``data.name``, mirroring the JS SDK)."""
+        data = self.raw.get("data")
+        if isinstance(data, dict) and data.get("name") is not None:
+            return str(data["name"])
+        return None
+
+    @property
+    def priority(self) -> int:
+        return int(self.raw.get("priority") or 0)
+
+    @property
+    def attempts(self) -> int:
+        return int(self.raw.get("attempts") or 0)
+
+    @property
+    def max_attempts(self) -> int:
+        return int(self.raw.get("maxAttempts") or 0)
+
+    @property
+    def created_at(self) -> Optional[int]:
+        value = self.raw.get("createdAt")
+        return int(value) if value is not None else None
+
+    @property
+    def started_at(self) -> Optional[int]:
+        value = self.raw.get("startedAt")
+        return int(value) if value is not None else None
+
+    @property
+    def completed_at(self) -> Optional[int]:
+        value = self.raw.get("completedAt")
+        return int(value) if value is not None else None
+
+    @property
+    def custom_id(self) -> Optional[str]:
+        value = self.raw.get("customId")
+        return str(value) if value is not None else None
+
+    @property
+    def tags(self) -> List[str]:
+        return list(self.raw.get("tags") or [])
+
+    @property
+    def group_id(self) -> Optional[str]:
+        value = self.raw.get("groupId")
+        return str(value) if value is not None else None
+
+    @property
+    def parent_id(self) -> Optional[str]:
+        value = self.raw.get("parentId")
+        return str(value) if value is not None else None
+
+    @property
+    def children_ids(self) -> List[str]:
+        return [str(c) for c in (self.raw.get("childrenIds") or [])]
+
+    @property
+    def depends_on(self) -> List[str]:
+        return [str(c) for c in (self.raw.get("dependsOn") or [])]
+
+    @property
+    def progress(self) -> int:
+        return int(self.raw.get("progress") or 0)
+
+    @property
+    def stacktrace(self) -> Optional[List[str]]:
+        value = self.raw.get("stacktrace")
+        return list(value) if value else None
+
+    @property
+    def state(self) -> Optional[str]:
+        value = self.raw.get("state")
+        return str(value) if value is not None else None
+
+    # ------------------------------------------------------------ operations
+
+    def _call(self, command: Dict[str, Any]) -> Dict[str, Any]:
+        if self._conn is None:
+            raise RuntimeError("Job is not bound to a connection")
+        return self._conn.call(_compact(command))
+
+    def update_progress(self, progress: int, message: Optional[str] = None) -> None:
+        """Report processing progress (0-100) to the server."""
+        self._call({"cmd": "Progress", "id": self.id, "progress": progress, "message": message})
+        if self._on_progress is not None:
+            self._on_progress(self, progress)
+
+    def log(self, message: str, level: Optional[str] = None) -> None:
+        """Append a log row to the job (AddLog)."""
+        self._call({"cmd": "AddLog", "id": self.id, "message": message, "level": level})
+
+    def heartbeat(self) -> None:
+        """Renew this job's lock (stall detection)."""
+        self._call({"cmd": "JobHeartbeat", "id": self.id, "token": self.token})
+
+    def extend_lock(self, duration_ms: int) -> None:
+        self._call({"cmd": "ExtendLock", "id": self.id, "token": self.token, "duration": duration_ms})
+
+    def get_state(self) -> str:
+        return str(self._call({"cmd": "GetState", "id": self.id}).get("state"))
+
+    def get_children_values(self) -> Dict[str, Any]:
+        response = self._call({"cmd": "GetChildrenValues", "id": self.id})
+        data = response.get("data")
+        if isinstance(data, dict) and "values" in data:
+            return dict(data.get("values") or {})
+        return dict(response.get("values") or {})
+
+    def remove(self) -> None:
+        self._call({"cmd": "Cancel", "id": self.id})
+
+    def discard(self) -> None:
+        self._call({"cmd": "Discard", "id": self.id})
+
+    def promote(self) -> None:
+        self._call({"cmd": "Promote", "id": self.id})
+
+    def retry(self) -> None:
+        """failed -> waiting (BullMQ contract; MoveToWait over TCP)."""
+        self._call({"cmd": "MoveToWait", "id": self.id})
+
+    def update_data(self, data: Any) -> None:
+        self._call({"cmd": "Update", "id": self.id, "data": data})
+
+    def change_priority(self, priority: int, lifo: Optional[bool] = None) -> None:
+        self._call({"cmd": "ChangePriority", "id": self.id, "priority": priority, "lifo": lifo})
+
+    def change_delay(self, delay_ms: int) -> None:
+        self._call({"cmd": "ChangeDelay", "id": self.id, "delay": delay_ms})
+
+    def move_to_delayed(self, delay_ms: int) -> None:
+        self._call({"cmd": "MoveToDelayed", "id": self.id, "delay": delay_ms})
+
+    def __repr__(self) -> str:  # pragma: no cover
+        return f"Job(id={self.id!r}, queue={self.queue!r}, name={self.name!r})"

bunqueue/options.py

@@ -0,0 +1,126 @@
+"""Job option mapping: pythonic kwargs -> wire fields (parity with TS client)."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Sequence, Union
+
+from .connection import _compact
+
+Backoff = Union[int, Dict[str, Any]]  # ms or {"type": "fixed"|"exponential", "delay", "maxDelay"}
+
+
+def job_payload(name: str, data: Any) -> Dict[str, Any]:
+    """Mirror the JS SDK: the job name travels inside ``data``."""
+    if data is None:
+        return {"name": name}
+    if isinstance(data, dict):
+        return {"name": name, **data}
+    return {"name": name, "payload": data}
+
+
+def build_cron_job_options(job_opts: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    """Map a cron template's Pythonic job options to the wire ``CronJobOptions``.
+
+    Mirrors the TS reference ``buildCronJobOptions`` (issue #86): the server
+    reads only camelCase keys (``maxAttempts``/``removeOnComplete``/…), so
+    SDK-native names (``attempts``/``remove_on_complete``/…) must be renamed
+    or the cron-spawned jobs silently fall back to JOB_DEFAULTS. Returns
+    ``None`` when nothing relevant is set (server keeps its own defaults).
+    """
+    if not job_opts:
+        return None
+    out: Dict[str, Any] = {}
+    if job_opts.get("attempts") is not None:
+        out["maxAttempts"] = job_opts["attempts"]
+    if job_opts.get("backoff") is not None:
+        out["backoff"] = job_opts["backoff"]
+    if job_opts.get("timeout") is not None:
+        out["timeout"] = job_opts["timeout"]
+    if job_opts.get("delay") is not None:
+        out["delay"] = job_opts["delay"]
+    if job_opts.get("stall_timeout") is not None:
+        out["stallTimeout"] = job_opts["stall_timeout"]
+    if isinstance(job_opts.get("remove_on_complete"), bool):
+        out["removeOnComplete"] = job_opts["remove_on_complete"]
+    if isinstance(job_opts.get("remove_on_fail"), bool):
+        out["removeOnFail"] = job_opts["remove_on_fail"]
+    return out or None
+
+
+def job_options(
+    *,
+    priority: Optional[int] = None,
+    delay: Optional[int] = None,
+    attempts: Optional[int] = None,
+    backoff: Optional[Backoff] = None,
+    ttl: Optional[int] = None,
+    timeout: Optional[int] = None,
+    job_id: Optional[str] = None,
+    unique_key: Optional[str] = None,
+    deduplication: Optional[Dict[str, Any]] = None,  # {"id", "ttl", "extend", "replace"}
+    depends_on: Optional[Sequence[str]] = None,
+    parent_id: Optional[str] = None,
+    children_ids: Optional[Sequence[str]] = None,
+    tags: Optional[Sequence[str]] = None,
+    group_id: Optional[str] = None,
+    lifo: Optional[bool] = None,
+    remove_on_complete: Optional[bool] = None,
+    remove_on_fail: Optional[bool] = None,
+    stall_timeout: Optional[int] = None,
+    durable: Optional[bool] = None,
+    repeat: Optional[Dict[str, Any]] = None,  # {"every", "limit", "pattern", ...} camelCase wire keys
+    debounce: Optional[Dict[str, Any]] = None,  # {"id", "ttl"}
+    stack_trace_limit: Optional[int] = None,
+    keep_logs: Optional[int] = None,
+    size_limit: Optional[int] = None,
+    timestamp: Optional[int] = None,
+    fail_parent_on_failure: Optional[bool] = None,
+    remove_dependency_on_failure: Optional[bool] = None,
+    ignore_dependency_on_failure: Optional[bool] = None,
+    continue_parent_on_failure: Optional[bool] = None,
+) -> Dict[str, Any]:
+    """Translate SDK options into the TCP PUSH field set (see buildPushPayload #88/#90)."""
+    dedup = None
+    if deduplication:
+        unique_key = unique_key or deduplication.get("id")
+        dedup = _compact(
+            {
+                "ttl": deduplication.get("ttl"),
+                "extend": deduplication.get("extend"),
+                "replace": deduplication.get("replace"),
+            }
+        ) or None
+    return _compact(
+        {
+            "priority": priority,
+            "delay": delay,
+            "maxAttempts": attempts,
+            "backoff": backoff,
+            "ttl": ttl,
+            "timeout": timeout,
+            "jobId": job_id,
+            "uniqueKey": unique_key,
+            "dedup": dedup,
+            "dependsOn": list(depends_on) if depends_on else None,
+            "parentId": parent_id,
+            "childrenIds": list(children_ids) if children_ids else None,
+            "tags": list(tags) if tags else None,
+            "groupId": group_id,
+            "lifo": lifo,
+            "removeOnComplete": remove_on_complete,
+            "removeOnFail": remove_on_fail,
+            "stallTimeout": stall_timeout,
+            "durable": durable,
+            "repeat": repeat,
+            "debounceId": debounce.get("id") if debounce else None,
+            "debounceTtl": debounce.get("ttl") if debounce else None,
+            "stackTraceLimit": stack_trace_limit,
+            "keepLogs": keep_logs,
+            "sizeLimit": size_limit,
+            "timestamp": timestamp,
+            "failParentOnFailure": fail_parent_on_failure,
+            "removeDependencyOnFailure": remove_dependency_on_failure,
+            "ignoreDependencyOnFailure": ignore_dependency_on_failure,
+            "continueParentOnFailure": continue_parent_on_failure,
+        }
+    )