labtasker-client 2.0.0__py3-none-any.whl

labtasker/client.py

@@ -0,0 +1,725 @@
+from __future__ import annotations
+
+import secrets
+import sys
+import time
+from collections.abc import Callable
+from pathlib import Path
+from typing import TypeVar
+
+import httpx
+from pydantic import TypeAdapter, ValidationError
+
+from labtasker.config import ResolvedConfig, resolve_config
+from labtasker.errors import APIError, TransportError
+from labtasker.local import (
+    ensure_local_server,
+    local_paths,
+    require_local_capabilities,
+    socket_transport,
+)
+from labtasker.models import (
+    BulkUpdateResult,
+    ClaimResponse,
+    CountResponse,
+    HealthResponse,
+    HeartbeatResponse,
+    Queue,
+    ResponseModel,
+    Task,
+    TaskPage,
+)
+from labtasker.types import JSONValue, TaskOrderField, TaskStatus, TaskUpdate
+from labtasker.validation import (
+    RequestValidationError,
+    validate_filter,
+    validate_identifier,
+    validate_int64,
+    validate_json_object,
+    validate_order_field,
+    validate_routes,
+    validate_run_id,
+    validate_status,
+    validate_task_id,
+    validate_task_name,
+    validate_task_update,
+    validate_unicode_scalar,
+)
+
+T = TypeVar("T")
+ModelT = TypeVar("ModelT", bound=ResponseModel)
+REQUEST_TIMEOUT_SECONDS = 10.0
+MAX_RETRY_ATTEMPTS = 3
+RETRY_BACKOFF_SECONDS = (0.05, 0.1)
+QUEUE_LIST_ADAPTER = TypeAdapter(list[Queue])
+
+
+class Client:
+    def __init__(
+        self,
+        url: str | None = None,
+        token: str | None = None,
+        queue: str | None = None,
+    ) -> None:
+        self._initialize(resolve_config(url=url, token=token, queue=queue))
+
+    @classmethod
+    def _from_local_directory(cls, directory: Path, *, queue: str) -> Client:
+        require_local_capabilities()
+        client = cls.__new__(cls)
+        client._initialize(
+            ResolvedConfig(
+                url=None,
+                queue=validate_identifier(queue, field="queue"),
+                token=None,
+                local=local_paths(directory),
+            )
+        )
+        return client
+
+    def _initialize(self, config: ResolvedConfig) -> None:
+        self._config = config
+        headers = {}
+        if self._config.token is not None:
+            headers["Authorization"] = f"Bearer {self._config.token}"
+        if self._config.local is None:
+            assert self._config.url is not None
+            self._http = httpx.Client(
+                base_url=f"{self._config.url}/api/v2/",
+                headers=headers,
+                timeout=REQUEST_TIMEOUT_SECONDS,
+            )
+        else:
+            self._http = httpx.Client(
+                base_url="http://labtasker/api/v2/",
+                transport=socket_transport(self._config.local),
+                timeout=REQUEST_TIMEOUT_SECONDS,
+            )
+        self._closed = False
+        self._endpoint_announced = False
+        self._local_ready = False
+
+    def __enter__(self) -> Client:
+        self._ensure_open()
+        return self
+
+    def __exit__(self, *_: object) -> None:
+        self.close()
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        self._http.close()
+
+    @property
+    def configuration(self) -> ResolvedConfig:
+        return self._config
+
+    def submit_task(
+        self,
+        args: dict[str, JSONValue] | None = None,
+        *,
+        name: str | None = None,
+        metadata: dict[str, JSONValue] | None = None,
+        priority: int = 0,
+        max_attempts: int = 3,
+        routes: list[str] | None = None,
+        task_id: str | None = None,
+        queue: str | None = None,
+    ) -> Task:
+        self._ensure_open()
+        normalized_args = validate_json_object({} if args is None else args, field="args")
+        normalized_metadata = validate_json_object(
+            {} if metadata is None else metadata,
+            field="metadata",
+        )
+        normalized_routes = validate_routes(["default"] if routes is None else routes)
+        body: dict[str, object] = {
+            "name": validate_task_name(name),
+            "args": normalized_args,
+            "metadata": normalized_metadata,
+            "priority": validate_int64(priority, field="priority"),
+            "max_attempts": validate_int64(
+                max_attempts,
+                field="max_attempts",
+                positive=True,
+            ),
+            "routes": normalized_routes,
+        }
+        queue_name = self._queue(queue)
+        if task_id is not None:
+            selected_id = validate_task_id(task_id)
+            return self._submit_with_id(queue_name, selected_id, body)
+
+        for _ in range(MAX_RETRY_ATTEMPTS):
+            selected_id = _generate_task_id()
+            try:
+                return self._submit_with_id(queue_name, selected_id, body)
+            except APIError as error:
+                if error.code != "task_id_conflict":
+                    raise
+        raise TransportError(
+            "Could not allocate a unique Task ID.",
+            {"operation": "submit_task", **self._operation_endpoint_details},
+        )
+
+    def get_task(self, task_id: str, *, queue: str | None = None) -> Task:
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        task_id = validate_task_id(task_id)
+        return self._call(
+            operation="get_task",
+            method="GET",
+            path=f"queues/{queue_name}/tasks/{task_id}",
+            parser=lambda response: _parse_model(response, Task, {200}),
+            retry=True,
+        )
+
+    def list_tasks(
+        self,
+        *,
+        status: TaskStatus | None = None,
+        name: str | None = None,
+        filter: str | None = None,
+        order_by: TaskOrderField = "created_at",
+        descending: bool = True,
+        limit: int = 100,
+        cursor: str | None = None,
+        queue: str | None = None,
+    ) -> TaskPage:
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        status = validate_status(status)
+        order_by = validate_order_field(order_by)
+        filter = validate_filter(filter)
+        if not isinstance(descending, bool):
+            raise RequestValidationError("descending must be a Boolean.")
+        if isinstance(limit, bool) or not isinstance(limit, int) or not 1 <= limit <= 1000:
+            raise RequestValidationError("limit must be an integer from 1 through 1000.")
+        if name is not None and not isinstance(name, str):
+            raise RequestValidationError("name selector must be a string or None.")
+        if cursor is not None and not isinstance(cursor, str):
+            raise RequestValidationError("cursor must be a string or None.")
+        params = _without_none(
+            {
+                "status": status,
+                "name": name,
+                "filter": filter,
+                "order_by": order_by,
+                "descending": "true" if descending else "false",
+                "limit": limit,
+                "cursor": cursor,
+            }
+        )
+        return self._call(
+            operation="list_tasks",
+            method="GET",
+            path=f"queues/{queue_name}/tasks",
+            params=params,
+            parser=lambda response: _parse_model(response, TaskPage, {200}),
+            retry=True,
+        )
+
+    def count_tasks(
+        self,
+        *,
+        status: TaskStatus | None = None,
+        name: str | None = None,
+        filter: str | None = None,
+        queue: str | None = None,
+    ) -> int:
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        status = validate_status(status)
+        filter = validate_filter(filter)
+        if name is not None and not isinstance(name, str):
+            raise RequestValidationError("name selector must be a string or None.")
+        result = self._call(
+            operation="count_tasks",
+            method="GET",
+            path=f"queues/{queue_name}/tasks/count",
+            params=_without_none({"status": status, "name": name, "filter": filter}),
+            parser=lambda response: _parse_model(response, CountResponse, {200}),
+            retry=True,
+        )
+        return result.count
+
+    def update_task(
+        self,
+        task_id: str,
+        changes: TaskUpdate,
+        *,
+        queue: str | None = None,
+    ) -> Task:
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        task_id = validate_task_id(task_id)
+        normalized = validate_task_update(changes)
+        return self._call(
+            operation="update_task",
+            method="PATCH",
+            path=f"queues/{queue_name}/tasks/{task_id}",
+            json=normalized,
+            parser=lambda response: _parse_model(response, Task, {200}),
+        )
+
+    def update_tasks(
+        self,
+        *,
+        filter: str,
+        changes: TaskUpdate,
+        queue: str | None = None,
+    ) -> BulkUpdateResult:
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        normalized_filter = validate_filter(filter, required=True)
+        normalized_changes = validate_task_update(changes)
+        return self._call(
+            operation="update_tasks",
+            method="PATCH",
+            path=f"queues/{queue_name}/tasks",
+            json={"filter": normalized_filter, "changes": normalized_changes},
+            parser=lambda response: _parse_model(response, BulkUpdateResult, {200}),
+        )
+
+    def cancel_task(self, task_id: str, *, queue: str | None = None) -> Task:
+        return self._task_action("cancel", task_id, queue=queue)
+
+    def requeue_task(self, task_id: str, *, queue: str | None = None) -> Task:
+        return self._task_action("requeue", task_id, queue=queue)
+
+    def delete_task(self, task_id: str, *, queue: str | None = None) -> None:
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        task_id = validate_task_id(task_id)
+        self._call(
+            operation="delete_task",
+            method="DELETE",
+            path=f"queues/{queue_name}/tasks/{task_id}",
+            parser=lambda response: _parse_none(response, {204}),
+        )
+
+    def create_queue(self, name: str) -> Queue:
+        self._ensure_open()
+        name = validate_identifier(name, field="queue")
+        return self._call(
+            operation="create_queue",
+            method="PUT",
+            path=f"queues/{name}",
+            parser=lambda response: _parse_model(response, Queue, {200, 201}),
+        )
+
+    def list_queues(self) -> list[Queue]:
+        self._ensure_open()
+        return self._call(
+            operation="list_queues",
+            method="GET",
+            path="queues",
+            parser=lambda response: _parse_queue_list(response, {200}),
+            retry=True,
+        )
+
+    def delete_queue(self, name: str, *, cascade: bool = False) -> None:
+        self._ensure_open()
+        name = validate_identifier(name, field="queue")
+        if not isinstance(cascade, bool):
+            raise RequestValidationError("cascade must be a Boolean.")
+        self._call(
+            operation="delete_queue",
+            method="DELETE",
+            path=f"queues/{name}",
+            params={"cascade": "true" if cascade else "false"},
+            parser=lambda response: _parse_none(response, {204}),
+        )
+
+    def _claim(
+        self,
+        *,
+        route: str,
+        run_id: str,
+        queue: str | None = None,
+    ) -> ClaimResponse | None:
+        """Claim one Task using a caller-generated idempotency token."""
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        normalized_route = validate_identifier(route, field="route")
+        normalized_run_id = validate_run_id(run_id)
+        return self._call(
+            operation="claim_task",
+            method="POST",
+            path=f"queues/{queue_name}/tasks/claim",
+            json={"route": normalized_route, "run_id": normalized_run_id},
+            parser=_parse_claim,
+            retry=True,
+        )
+
+    def _health(self) -> HealthResponse:
+        """Validate the deployment health and v2 protocol before Worker claims."""
+        return self._call(
+            operation="worker_health",
+            method="GET",
+            path=(
+                f"{self._config.url}/health"
+                if self._config.url is not None
+                else "http://labtasker/health"
+            ),
+            parser=lambda response: _parse_model(response, HealthResponse, {200}),
+            retry=True,
+        )
+
+    def _heartbeat(
+        self,
+        *,
+        task_id: str,
+        run_id: str,
+        queue: str | None = None,
+    ) -> HeartbeatResponse:
+        return self._run_action(
+            "heartbeat",
+            task_id=task_id,
+            run_id=run_id,
+            queue=queue,
+            body={},
+            parser=lambda response: _parse_model(response, HeartbeatResponse, {200}),
+        )
+
+    def _complete(
+        self,
+        *,
+        task_id: str,
+        run_id: str,
+        result: dict[str, JSONValue],
+        queue: str | None = None,
+    ) -> None:
+        normalized = validate_json_object(result, field="result")
+        self._run_action(
+            "complete",
+            task_id=task_id,
+            run_id=run_id,
+            queue=queue,
+            body={"result": normalized},
+            parser=lambda response: _parse_none(response, {204}),
+        )
+
+    def _fail(
+        self,
+        *,
+        task_id: str,
+        run_id: str,
+        error_type: str,
+        message: str,
+        traceback: str | None,
+        queue: str | None = None,
+    ) -> None:
+        error = {
+            "type": validate_unicode_scalar(error_type, field="error.type"),
+            "message": validate_unicode_scalar(message, field="error.message"),
+            "traceback": (
+                None
+                if traceback is None
+                else validate_unicode_scalar(traceback, field="error.traceback")
+            ),
+        }
+        self._run_action(
+            "fail",
+            task_id=task_id,
+            run_id=run_id,
+            queue=queue,
+            body={"error": error},
+            parser=lambda response: _parse_none(response, {204}),
+        )
+
+    def _unclaim(
+        self,
+        *,
+        task_id: str,
+        run_id: str,
+        queue: str | None = None,
+    ) -> None:
+        self._run_action(
+            "unclaim",
+            task_id=task_id,
+            run_id=run_id,
+            queue=queue,
+            body={},
+            parser=lambda response: _parse_none(response, {204}),
+        )
+
+    def _submit_with_id(
+        self,
+        queue: str,
+        task_id: str,
+        body: dict[str, object],
+    ) -> Task:
+        return self._call(
+            operation="submit_task",
+            method="PUT",
+            path=f"queues/{queue}/tasks/{task_id}",
+            json=body,
+            parser=lambda response: _parse_model(response, Task, {200, 201}),
+            retry=True,
+        )
+
+    def _task_action(self, action: str, task_id: str, *, queue: str | None) -> Task:
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        task_id = validate_task_id(task_id)
+        return self._call(
+            operation=f"{action}_task",
+            method="POST",
+            path=f"queues/{queue_name}/tasks/{task_id}/{action}",
+            parser=lambda response: _parse_model(response, Task, {200}),
+        )
+
+    def _run_action(
+        self,
+        action: str,
+        *,
+        task_id: str,
+        run_id: str,
+        queue: str | None,
+        body: dict[str, object],
+        parser: Callable[[httpx.Response], T],
+    ) -> T:
+        self._ensure_open()
+        queue_name = self._queue(queue)
+        normalized_task_id = validate_task_id(task_id)
+        normalized_run_id = validate_run_id(run_id)
+        return self._call(
+            operation=f"{action}_run",
+            method="POST",
+            path=f"queues/{queue_name}/tasks/{normalized_task_id}/{action}",
+            json={"run_id": normalized_run_id, **body},
+            parser=parser,
+        )
+
+    def _queue(self, queue: str | None) -> str:
+        return validate_identifier(self._config.queue if queue is None else queue, field="queue")
+
+    def _call(
+        self,
+        *,
+        operation: str,
+        method: str,
+        path: str,
+        parser: Callable[[httpx.Response], T],
+        json: object | None = None,
+        params: dict[str, str | int] | None = None,
+        retry: bool = False,
+    ) -> T:
+        self._ensure_open()
+        self._prepare_endpoint()
+        attempts = MAX_RETRY_ATTEMPTS if retry else 1
+        last_transport_error: TransportError | None = None
+        local_connect_recovery_used = False
+        attempt = 0
+        while attempt < attempts:
+            try:
+                response = self._http.request(method, path, json=json, params=params)
+            except httpx.RequestError as error:
+                last_transport_error = self._connection_error(operation)
+                can_recover_local_connect = (
+                    self._config.local is not None
+                    and isinstance(error, (httpx.ConnectError, httpx.ConnectTimeout))
+                    and not local_connect_recovery_used
+                )
+                if can_recover_local_connect:
+                    local_connect_recovery_used = True
+                    self._local_ready = False
+                    self._ensure_local_available()
+                    if attempt + 1 == attempts:
+                        attempts += 1
+                elif attempt + 1 == attempts:
+                    raise last_transport_error from error
+            else:
+                self._announce_http_endpoint()
+                if response.is_error:
+                    try:
+                        api_error = _parse_api_error(response)
+                    except TransportError as error:
+                        last_transport_error = _with_operation(
+                            error,
+                            operation,
+                            self._operation_endpoint_details,
+                        )
+                        if attempt + 1 == attempts:
+                            raise last_transport_error from error
+                        _backoff(attempt)
+                        attempt += 1
+                        continue
+                    if retry and api_error.code == "database_busy" and attempt + 1 < attempts:
+                        _backoff(attempt)
+                        attempt += 1
+                        continue
+                    raise api_error
+                try:
+                    return parser(response)
+                except TransportError as error:
+                    last_transport_error = _with_operation(
+                        error, operation, self._operation_endpoint_details
+                    )
+                    if attempt + 1 == attempts:
+                        raise last_transport_error from error
+            _backoff(attempt)
+            attempt += 1
+        if last_transport_error is None:
+            raise AssertionError("Request loop ended without a result or error.")
+        raise last_transport_error
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("Client is closed.")
+
+    @property
+    def _operation_endpoint_details(self) -> dict[str, object]:
+        if self._config.local is None:
+            return {"url": self._config.url}
+        return {
+            "directory": str(self._config.local.directory),
+            "socket": str(self._config.local.socket),
+        }
+
+    def _prepare_endpoint(self) -> None:
+        if self._config.local is not None and not self._local_ready:
+            self._ensure_local_available()
+
+    def _announce_http_endpoint(self) -> None:
+        if self._endpoint_announced or self._config.local is not None:
+            return
+        assert self._config.url is not None
+        transport = self._config.url.partition(":")[0]
+        print(
+            f"[labtasker] connected server=remote transport={transport} url={self._config.url}",
+            file=sys.stderr,
+        )
+        self._endpoint_announced = True
+
+    def _ensure_local_available(self) -> None:
+        paths = self._config.local
+        if paths is None:
+            return
+        result = ensure_local_server(paths, emit=self._emit_local_transition)
+        pid = result.pid if result.pid is not None else "unknown"
+        version = result.server_version if result.server_version is not None else "unknown"
+        print(
+            "[labtasker] connected server=local transport=unix "
+            f"directory={paths.directory} database={paths.database} socket={paths.socket} "
+            f"pid={pid} version={version}",
+            file=sys.stderr,
+        )
+        self._endpoint_announced = True
+        self._local_ready = True
+
+    @staticmethod
+    def _emit_local_transition(message: str) -> None:
+        print(f"[labtasker] {message}", file=sys.stderr)
+
+    def _connection_error(self, operation: str) -> TransportError:
+        details: dict[str, object] = {"operation": operation}
+        if self._config.local is None:
+            details["url"] = self._config.url
+        else:
+            details.update(
+                {
+                    "state": "unhealthy",
+                    "directory": str(self._config.local.directory),
+                    "database": str(self._config.local.database),
+                    "socket": str(self._config.local.socket),
+                    "log": str(self._config.local.log),
+                }
+            )
+        return TransportError("The Labtasker Server could not be reached.", details)
+
+
+def _parse_model(
+    response: httpx.Response,
+    model: type[ModelT],
+    statuses: set[int],
+) -> ModelT:
+    _require_status(response, statuses)
+    try:
+        return model.model_validate_json(response.content, strict=True)
+    except (ValidationError, ValueError, TypeError) as error:
+        raise TransportError(
+            "The Server returned an invalid success response.",
+            {"http_status": response.status_code},
+        ) from error
+
+
+def _parse_queue_list(response: httpx.Response, statuses: set[int]) -> list[Queue]:
+    _require_status(response, statuses)
+    try:
+        return QUEUE_LIST_ADAPTER.validate_json(response.content, strict=True)
+    except (ValidationError, ValueError, TypeError) as error:
+        raise TransportError(
+            "The Server returned an invalid Queue list.",
+            {"http_status": response.status_code},
+        ) from error
+
+
+def _parse_none(response: httpx.Response, statuses: set[int]) -> None:
+    _require_status(response, statuses)
+    if response.content:
+        raise TransportError(
+            "The Server returned an unexpected response body.",
+            {"http_status": response.status_code},
+        )
+
+
+def _parse_claim(response: httpx.Response) -> ClaimResponse | None:
+    if response.status_code == 204:
+        _parse_none(response, {204})
+        return None
+    return _parse_model(response, ClaimResponse, {200})
+
+
+def _require_status(response: httpx.Response, statuses: set[int]) -> None:
+    if response.status_code not in statuses:
+        raise TransportError(
+            "The Server returned an unexpected success status.",
+            {"http_status": response.status_code},
+        )
+
+
+def _parse_api_error(response: httpx.Response) -> APIError:
+    try:
+        payload = response.json()
+        if not isinstance(payload, dict) or set(payload) != {"error"}:
+            raise ValueError
+        error = payload["error"]
+        if not isinstance(error, dict) or set(error) != {"code", "message", "details"}:
+            raise ValueError
+        code = error["code"]
+        message = error["message"]
+        details = error["details"]
+        if not isinstance(code, str) or not isinstance(message, str):
+            raise ValueError
+        normalized_details = validate_json_object(details, field="error.details")
+    except (ValueError, TypeError) as error:
+        raise TransportError(
+            "The Server returned an invalid error response.",
+            {"http_status": response.status_code},
+        ) from error
+    return APIError(response.status_code, code, message, normalized_details)
+
+
+def _with_operation(
+    error: TransportError,
+    operation: str,
+    endpoint: dict[str, object],
+) -> TransportError:
+    return TransportError(error.message, {**error.details, "operation": operation, **endpoint})
+
+
+def _generate_task_id() -> str:
+    return f"t_{secrets.token_urlsafe(9)}"
+
+
+def _without_none(values: dict[str, str | int | None]) -> dict[str, str | int]:
+    return {key: value for key, value in values.items() if value is not None}
+
+
+def _backoff(attempt: int) -> None:
+    if attempt < len(RETRY_BACKOFF_SECONDS):
+        time.sleep(RETRY_BACKOFF_SECONDS[attempt])