labtasker-client 2.0.0__py3-none-any.whl

labtasker/models.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, field_validator
+
+from labtasker.types import JSONValue, TaskStatus
+from labtasker.validation import (
+    validate_identifier,
+    validate_int64,
+    validate_json_object,
+    validate_routes,
+    validate_run_id,
+    validate_task_id,
+    validate_task_name,
+    validate_unicode_scalar,
+)
+
+
+class ResponseModel(BaseModel):
+    model_config = ConfigDict(extra="ignore", frozen=True, strict=True)
+
+
+class LastError(ResponseModel):
+    type: str
+    message: str
+    traceback: str | None
+    occurred_at: datetime
+    attempt: int
+    run_id: str
+
+    @field_validator("type", "message", "traceback")
+    @classmethod
+    def validate_strings(cls, value: str | None, info: object) -> str | None:
+        if value is None:
+            return None
+        return validate_unicode_scalar(value, field=getattr(info, "field_name", "last_error"))
+
+    @field_validator("occurred_at")
+    @classmethod
+    def validate_time(cls, value: datetime) -> datetime:
+        return _utc_datetime(value)
+
+    @field_validator("attempt")
+    @classmethod
+    def validate_attempt(cls, value: int) -> int:
+        return validate_int64(value, field="attempt")
+
+    @field_validator("run_id")
+    @classmethod
+    def validate_run(cls, value: str) -> str:
+        return validate_run_id(value)
+
+
+class Task(ResponseModel):
+    id: str
+    queue: str
+    status: TaskStatus
+    name: str | None
+    args: dict[str, JSONValue]
+    metadata: dict[str, JSONValue]
+    priority: int
+    attempt: int
+    max_attempts: int
+    routes: list[str]
+    result: dict[str, JSONValue]
+    last_error: LastError | None
+    last_route: str | None
+    created_at: datetime
+    updated_at: datetime
+    started_at: datetime | None
+    finished_at: datetime | None
+
+    @field_validator("id")
+    @classmethod
+    def validate_id(cls, value: str) -> str:
+        return validate_task_id(value)
+
+    @field_validator("queue")
+    @classmethod
+    def validate_queue(cls, value: str) -> str:
+        return validate_identifier(value, field="queue")
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, value: str | None) -> str | None:
+        return validate_task_name(value)
+
+    @field_validator("args", "metadata", "result")
+    @classmethod
+    def validate_objects(cls, value: dict[str, JSONValue], info: object) -> dict[str, JSONValue]:
+        return validate_json_object(value, field=getattr(info, "field_name", "task"))
+
+    @field_validator("priority", "attempt")
+    @classmethod
+    def validate_numbers(cls, value: int, info: object) -> int:
+        return validate_int64(value, field=getattr(info, "field_name", "task"))
+
+    @field_validator("max_attempts")
+    @classmethod
+    def validate_max_attempts(cls, value: int) -> int:
+        return validate_int64(value, field="max_attempts", positive=True)
+
+    @field_validator("routes")
+    @classmethod
+    def validate_task_routes(cls, value: list[str]) -> list[str]:
+        routes = validate_routes(value)
+        if value != routes:
+            raise ValueError("routes must be sorted lexicographically")
+        return routes
+
+    @field_validator("last_route")
+    @classmethod
+    def validate_last_route(cls, value: str | None) -> str | None:
+        return None if value is None else validate_identifier(value, field="last_route")
+
+    @field_validator("created_at", "updated_at", "started_at", "finished_at")
+    @classmethod
+    def validate_times(cls, value: datetime | None) -> datetime | None:
+        return None if value is None else _utc_datetime(value)
+
+
+class TaskInfo(Task):
+    run_id: str
+    run_dir: Path
+
+    @field_validator("run_id")
+    @classmethod
+    def validate_run(cls, value: str) -> str:
+        return validate_run_id(value)
+
+    @field_validator("run_dir")
+    @classmethod
+    def validate_run_dir(cls, value: Path) -> Path:
+        if not value.is_absolute():
+            raise ValueError("run_dir must be absolute")
+        return value
+
+
+class TaskPage(ResponseModel):
+    items: list[Task]
+    next_cursor: str | None
+
+
+class ClaimResponse(ResponseModel):
+    task: Task
+    run_id: str
+    lease_expires_at: datetime
+
+    @field_validator("run_id")
+    @classmethod
+    def validate_run(cls, value: str) -> str:
+        return validate_run_id(value)
+
+    @field_validator("lease_expires_at")
+    @classmethod
+    def validate_lease_expires_at(cls, value: datetime) -> datetime:
+        return _utc_datetime(value)
+
+
+class HeartbeatResponse(ResponseModel):
+    lease_expires_at: datetime
+
+    @field_validator("lease_expires_at")
+    @classmethod
+    def validate_lease_expires_at(cls, value: datetime) -> datetime:
+        return _utc_datetime(value)
+
+
+class HealthResponse(ResponseModel):
+    status: Literal["ok"]
+    api_version: Literal["2"]
+    database: Literal["ok"]
+
+
+class Queue(ResponseModel):
+    name: str
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, value: str) -> str:
+        return validate_identifier(value, field="queue")
+
+
+class BulkUpdateResult(ResponseModel):
+    matched: int
+    updated: int
+
+    @field_validator("matched", "updated")
+    @classmethod
+    def validate_count(cls, value: int, info: object) -> int:
+        value = validate_int64(value, field=getattr(info, "field_name", "count"))
+        if value < 0:
+            raise ValueError("count must be non-negative")
+        return value
+
+
+class CountResponse(ResponseModel):
+    count: int
+
+    @field_validator("count")
+    @classmethod
+    def validate_count(cls, value: int) -> int:
+        value = validate_int64(value, field="count")
+        if value < 0:
+            raise ValueError("count must be non-negative")
+        return value
+
+
+def _utc_datetime(value: datetime) -> datetime:
+    offset = value.utcoffset()
+    if value.tzinfo is None or offset is None:
+        raise ValueError("timestamp must be timezone-aware")
+    if offset.total_seconds() != 0:
+        raise ValueError("timestamp must use UTC")
+    return value.astimezone(UTC)

labtasker/paths.py

@@ -0,0 +1,34 @@
+"""Shared object-only dot paths used by Python and command Workers."""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+_PATH_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*$")
+
+
+class PathError(ValueError):
+    pass
+
+
+def parse_path(value: str) -> tuple[str, ...]:
+    if not isinstance(value, str) or not _PATH_RE.fullmatch(value):
+        raise PathError(
+            "path must contain dot-separated ASCII identifiers matching [A-Za-z_][A-Za-z0-9_]*"
+        )
+    return tuple(value.split("."))
+
+
+def select_path(value: object, path: tuple[str, ...]) -> Any:
+    current = value
+    traversed: list[str] = []
+    for segment in path:
+        traversed.append(segment)
+        if not isinstance(current, dict):
+            parent = ".".join(traversed[:-1]) or "<root>"
+            raise PathError(f"path {'.'.join(path)!r} cannot traverse non-object {parent!r}")
+        if segment not in current:
+            raise PathError(f"path {'.'.join(path)!r} is missing key {segment!r}")
+        current = current[segment]
+    return current

labtasker/py.typed

@@ -0,0 +1 @@
+

labtasker/tee.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import threading
+import time
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from time import struct_time
+from typing import TextIO, cast
+
+_ACTIVE_TEE: WorkerTee | None = None
+_FORK_HOOK_INSTALLED = False
+
+
+class _UTCFormatter(logging.Formatter):
+    converter: Callable[[float | None], struct_time] = time.gmtime
+
+
+class _TeeStream:
+    def __init__(self, original: TextIO, lock: threading.RLock) -> None:
+        self._original = original
+        self._lock = lock
+        self._destination: TextIO | None = None
+
+    def write(self, value: str) -> int:
+        with self._lock:
+            written = self._original.write(value)
+            if self._destination is not None:
+                self._destination.write(value)
+            return written
+
+    def flush(self) -> None:
+        with self._lock:
+            self._original.flush()
+            if self._destination is not None:
+                self._destination.flush()
+
+    def set_destination(self, destination: TextIO | None) -> None:
+        with self._lock:
+            self._destination = destination
+
+    def __getattr__(self, name: str) -> object:
+        return getattr(self._original, name)
+
+
+class WorkerTee:
+    def __init__(self) -> None:
+        self._lock = threading.RLock()
+        self._stdout: _TeeStream | None = None
+        self._stderr: _TeeStream | None = None
+        self._original_stdout: TextIO | None = None
+        self._original_stderr: TextIO | None = None
+        self._destination: TextIO | None = None
+
+    def __enter__(self) -> WorkerTee:
+        global _ACTIVE_TEE, _FORK_HOOK_INSTALLED
+        self._original_stdout = sys.stdout
+        self._original_stderr = sys.stderr
+        self._stdout = _TeeStream(sys.stdout, self._lock)
+        self._stderr = _TeeStream(sys.stderr, self._lock)
+        sys.stdout = cast(TextIO, self._stdout)
+        sys.stderr = cast(TextIO, self._stderr)
+        if not _FORK_HOOK_INSTALLED and hasattr(os, "register_at_fork"):
+            os.register_at_fork(after_in_child=_clear_tee_after_fork)
+            _FORK_HOOK_INSTALLED = True
+        _ACTIVE_TEE = self
+        return self
+
+    def __exit__(self, *_: object) -> None:
+        global _ACTIVE_TEE
+        self.clear_destination()
+        if self._original_stdout is not None:
+            sys.stdout = self._original_stdout
+        if self._original_stderr is not None:
+            sys.stderr = self._original_stderr
+        if _ACTIVE_TEE is self:
+            _ACTIVE_TEE = None
+
+    @contextmanager
+    def capture(self, path: Path) -> Iterator[None]:
+        if self._destination is not None:
+            raise RuntimeError("A Worker log destination is already active.")
+        with path.open("a", encoding="utf-8", errors="backslashreplace") as destination:
+            self._destination = destination
+            if self._stdout is not None:
+                self._stdout.set_destination(destination)
+            if self._stderr is not None:
+                self._stderr.set_destination(destination)
+            try:
+                yield
+            finally:
+                self.clear_destination()
+
+    def clear_destination(self) -> None:
+        if self._stdout is not None:
+            self._stdout.set_destination(None)
+        if self._stderr is not None:
+            self._stderr.set_destination(None)
+        if self._destination is not None:
+            self._destination.flush()
+            self._destination = None
+
+
+def configure_worker_logger() -> logging.Logger:
+    logger = logging.getLogger("labtasker")
+    if not logger.hasHandlers():
+        handler = logging.StreamHandler()
+        handler.setLevel(logging.INFO)
+        handler.setFormatter(_worker_log_formatter())
+        logger.addHandler(handler)
+        if logger.level == logging.NOTSET:
+            logger.setLevel(logging.INFO)
+    return logger
+
+
+def _worker_log_formatter() -> logging.Formatter:
+    return _UTCFormatter(
+        "%(asctime)s.%(msecs)03dZ %(levelname)s [labtasker] %(message)s",
+        datefmt="%Y-%m-%dT%H:%M:%S",
+    )
+
+
+def _clear_tee_after_fork() -> None:
+    if _ACTIVE_TEE is not None:
+        _ACTIVE_TEE.clear_destination()

labtasker/types.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import Literal, TypeAlias, TypedDict
+
+from pydantic import JsonValue as PydanticJSONValue
+
+JSONValue: TypeAlias = PydanticJSONValue
+TaskStatus: TypeAlias = Literal["pending", "running", "succeeded", "failed", "cancelled"]
+TaskOrderField: TypeAlias = Literal[
+    "id",
+    "name",
+    "status",
+    "priority",
+    "attempt",
+    "max_attempts",
+    "last_route",
+    "created_at",
+    "updated_at",
+    "started_at",
+    "finished_at",
+]
+
+
+class TaskUpdate(TypedDict, total=False):
+    name: str | None
+    args: dict[str, JSONValue]
+    metadata: dict[str, JSONValue]
+    priority: int
+    max_attempts: int
+    routes: list[str]
+    result: dict[str, JSONValue]

labtasker/validation.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import math
+import re
+import unicodedata
+from collections.abc import Sequence
+from typing import cast
+
+from labtasker.errors import ConfigError
+from labtasker.types import JSONValue, TaskOrderField, TaskStatus, TaskUpdate
+
+INT64_MIN = -(2**63)
+INT64_MAX = 2**63 - 1
+MAX_JSON_DEPTH = 64
+MAX_FILTER_BYTES = 8192
+IDENTIFIER_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")
+TASK_ID_RE = re.compile(r"^t_[A-Za-z0-9_-]{12}$")
+RUN_ID_RE = re.compile(r"^r_[A-Za-z0-9_-]{12}$")
+TASK_STATUSES = {"pending", "running", "succeeded", "failed", "cancelled"}
+TASK_ORDER_FIELDS = {
+    "id",
+    "name",
+    "status",
+    "priority",
+    "attempt",
+    "max_attempts",
+    "last_route",
+    "created_at",
+    "updated_at",
+    "started_at",
+    "finished_at",
+}
+
+
+class RequestValidationError(ValueError):
+    pass
+
+
+def validate_json_object(value: object, *, field: str) -> dict[str, JSONValue]:
+    if not isinstance(value, dict):
+        raise RequestValidationError(f"{field} must be a JSON object.")
+    validate_json_value(value, field=field)
+    return cast(dict[str, JSONValue], value)
+
+
+def validate_json_value(value: object, *, field: str) -> None:
+    active_containers: set[int] = set()
+
+    def walk(current: object, depth: int, path: str) -> None:
+        if isinstance(current, str):
+            validate_unicode_scalar(current, field=path)
+            return
+        if current is None or isinstance(current, bool):
+            return
+        if isinstance(current, int):
+            if not INT64_MIN <= current <= INT64_MAX:
+                raise RequestValidationError(f"{path} is outside the signed 64-bit range.")
+            return
+        if isinstance(current, float):
+            if not math.isfinite(current):
+                raise RequestValidationError(f"{path} must be finite.")
+            return
+        if isinstance(current, dict):
+            if depth >= MAX_JSON_DEPTH:
+                raise RequestValidationError(f"{field} exceeds JSON depth {MAX_JSON_DEPTH}.")
+            identity = id(current)
+            if identity in active_containers:
+                raise RequestValidationError(f"{field} contains a cycle.")
+            active_containers.add(identity)
+            try:
+                for key, child in current.items():
+                    if not isinstance(key, str):
+                        raise RequestValidationError(f"{path} has a non-string object key.")
+                    validate_unicode_scalar(key, field=f"{path}.<key>")
+                    walk(child, depth + 1, f"{path}.{key}")
+            finally:
+                active_containers.remove(identity)
+            return
+        if isinstance(current, Sequence) and not isinstance(current, (str, bytes, bytearray)):
+            if not isinstance(current, list):
+                raise RequestValidationError(f"{path} must use JSON arrays, not Python sequences.")
+            if depth >= MAX_JSON_DEPTH:
+                raise RequestValidationError(f"{field} exceeds JSON depth {MAX_JSON_DEPTH}.")
+            identity = id(current)
+            if identity in active_containers:
+                raise RequestValidationError(f"{field} contains a cycle.")
+            active_containers.add(identity)
+            try:
+                for index, child in enumerate(current):
+                    walk(child, depth + 1, f"{path}[{index}]")
+            finally:
+                active_containers.remove(identity)
+            return
+        raise RequestValidationError(f"{path} is not representable in strict JSON.")
+
+    walk(value, 0, field)
+
+
+def validate_unicode_scalar(value: str, *, field: str) -> str:
+    if any(0xD800 <= ord(character) <= 0xDFFF for character in value):
+        raise RequestValidationError(f"{field} contains a lone Unicode surrogate.")
+    return value
+
+
+def validate_identifier(value: object, *, field: str) -> str:
+    if not isinstance(value, str) or not IDENTIFIER_RE.fullmatch(value):
+        raise RequestValidationError(f"{field} must match [A-Za-z0-9][A-Za-z0-9._-]{{0,127}}.")
+    return value
+
+
+def validate_task_id(value: object) -> str:
+    if not isinstance(value, str) or not TASK_ID_RE.fullmatch(value):
+        raise RequestValidationError("task_id must match t_[A-Za-z0-9_-]{12}.")
+    return value
+
+
+def validate_run_id(value: object) -> str:
+    if not isinstance(value, str) or not RUN_ID_RE.fullmatch(value):
+        raise RequestValidationError("run_id must match r_[A-Za-z0-9_-]{12}.")
+    return value
+
+
+def validate_task_name(value: object) -> str | None:
+    if value is None:
+        return None
+    if not isinstance(value, str):
+        raise RequestValidationError("name must be a string or None.")
+    validate_unicode_scalar(value, field="name")
+    if len(value) > 256:
+        raise RequestValidationError("name exceeds 256 Unicode code points.")
+    if any(unicodedata.category(character) == "Cc" for character in value):
+        raise RequestValidationError("name contains a control character.")
+    return value
+
+
+def validate_int64(value: object, *, field: str, positive: bool = False) -> int:
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise RequestValidationError(f"{field} must be an integer.")
+    if not INT64_MIN <= value <= INT64_MAX or (positive and value <= 0):
+        qualifier = "a positive signed 64-bit integer" if positive else "a signed 64-bit integer"
+        raise RequestValidationError(f"{field} must be {qualifier}.")
+    return value
+
+
+def validate_routes(value: object) -> list[str]:
+    if not isinstance(value, list) or not value:
+        raise RequestValidationError("routes must be a non-empty list of strings.")
+    routes = [validate_identifier(route, field="route") for route in value]
+    if len(routes) != len(set(routes)):
+        raise RequestValidationError("routes must not contain duplicates.")
+    return sorted(routes)
+
+
+def validate_filter(value: object, *, required: bool = False) -> str | None:
+    if value is None:
+        if required:
+            raise RequestValidationError("filter must be a non-empty string.")
+        return None
+    if not isinstance(value, str) or not value.strip():
+        raise RequestValidationError("filter must be a non-empty string.")
+    validate_unicode_scalar(value, field="filter")
+    if len(value.encode("utf-8")) > MAX_FILTER_BYTES:
+        raise RequestValidationError(f"filter exceeds {MAX_FILTER_BYTES} bytes.")
+    return value
+
+
+def validate_status(value: object | None) -> TaskStatus | None:
+    if value is None:
+        return None
+    if not isinstance(value, str) or value not in TASK_STATUSES:
+        raise RequestValidationError("status is not a valid Task status.")
+    return cast(TaskStatus, value)
+
+
+def validate_order_field(value: object) -> TaskOrderField:
+    if not isinstance(value, str) or value not in TASK_ORDER_FIELDS:
+        raise RequestValidationError("order_by is not a supported Task field.")
+    return cast(TaskOrderField, value)
+
+
+def validate_task_update(changes: object) -> TaskUpdate:
+    if not isinstance(changes, dict) or not changes:
+        raise RequestValidationError("changes must be a non-empty object.")
+    allowed = {"name", "args", "metadata", "priority", "max_attempts", "routes", "result"}
+    unknown = set(changes) - allowed
+    if unknown:
+        raise RequestValidationError(f"changes contains unsupported fields: {sorted(unknown)!r}.")
+    normalized: dict[str, object] = {}
+    for field, value in changes.items():
+        if field == "name":
+            normalized[field] = validate_task_name(value)
+        elif field in {"args", "metadata", "result"}:
+            normalized[field] = validate_json_object(value, field=field)
+        elif field == "priority":
+            normalized[field] = validate_int64(value, field=field)
+        elif field == "max_attempts":
+            normalized[field] = validate_int64(value, field=field, positive=True)
+        elif field == "routes":
+            normalized[field] = validate_routes(value)
+    return cast(TaskUpdate, normalized)
+
+
+def invalid_config(message: str, *, source: str, field: str | None = None) -> ConfigError:
+    details = {"source": source}
+    if field is not None:
+        details["field"] = field
+    return ConfigError("invalid_config", message, details)