labtasker-server 2.0.0__py3-none-any.whl

labtasker_server/migrations/versions/0001_initial.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+revision: str = "0001_initial"
+down_revision: str | None = None
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    op.create_table("queues", sa.Column("name", sa.String(length=128), primary_key=True))
+    op.create_table(
+        "tasks",
+        sa.Column("queue_name", sa.String(length=128), nullable=False),
+        sa.Column("task_id", sa.String(length=14), nullable=False),
+        sa.Column("status", sa.String(length=16), nullable=False),
+        sa.Column("name", sa.String(length=256)),
+        sa.Column("args_json", sa.Text(), nullable=False),
+        sa.Column("metadata_json", sa.Text(), nullable=False),
+        sa.Column("result_json", sa.Text(), nullable=False),
+        sa.Column("priority", sa.Integer(), nullable=False),
+        sa.Column("attempt", sa.Integer(), nullable=False),
+        sa.Column("max_attempts", sa.Integer(), nullable=False),
+        sa.Column("created_at_us", sa.Integer(), nullable=False),
+        sa.Column("updated_at_us", sa.Integer(), nullable=False),
+        sa.Column("last_route", sa.String(length=128)),
+        sa.Column("started_at_us", sa.Integer()),
+        sa.Column("finished_at_us", sa.Integer()),
+        sa.Column("last_error_json", sa.Text()),
+        sa.Column("creation_hash", sa.String(length=64), nullable=False),
+        sa.Column("active_run_id", sa.String(length=14)),
+        sa.Column("lease_expires_at_us", sa.Integer()),
+        sa.Column("last_terminal_run_id", sa.String(length=14)),
+        sa.Column("last_terminal_action", sa.String(length=32)),
+        sa.Column("pending_at_us", sa.Integer()),
+        sa.CheckConstraint(
+            "status IN ('pending','running','succeeded','failed','cancelled')",
+            name="ck_tasks_status",
+        ),
+        sa.CheckConstraint("attempt >= 0", name="ck_tasks_attempt_nonnegative"),
+        sa.CheckConstraint("max_attempts > 0", name="ck_tasks_max_attempts_positive"),
+        sa.CheckConstraint(
+            "last_terminal_action IS NULL OR last_terminal_action IN "
+            "('complete','fail','unclaim','heartbeat_expired','cancel')",
+            name="ck_tasks_terminal_action",
+        ),
+        sa.CheckConstraint(
+            "json_valid(args_json) AND json_type(args_json) = 'object'",
+            name="ck_args_json",
+        ),
+        sa.CheckConstraint(
+            "json_valid(metadata_json) AND json_type(metadata_json) = 'object'",
+            name="ck_metadata_json",
+        ),
+        sa.CheckConstraint(
+            "json_valid(result_json) AND json_type(result_json) = 'object'",
+            name="ck_result_json",
+        ),
+        sa.CheckConstraint(
+            "(status = 'pending' AND pending_at_us IS NOT NULL "
+            "AND active_run_id IS NULL AND lease_expires_at_us IS NULL "
+            "AND attempt < max_attempts) OR status != 'pending'",
+            name="ck_tasks_pending_state",
+        ),
+        sa.CheckConstraint(
+            "(status = 'running' AND pending_at_us IS NULL "
+            "AND active_run_id IS NOT NULL AND lease_expires_at_us IS NOT NULL) "
+            "OR status != 'running'",
+            name="ck_tasks_running_state",
+        ),
+        sa.CheckConstraint(
+            "(status IN ('succeeded','failed','cancelled') AND pending_at_us IS NULL "
+            "AND active_run_id IS NULL AND lease_expires_at_us IS NULL) "
+            "OR status NOT IN ('succeeded','failed','cancelled')",
+            name="ck_tasks_terminal_state",
+        ),
+        sa.ForeignKeyConstraint(["queue_name"], ["queues.name"], ondelete="CASCADE"),
+        sa.PrimaryKeyConstraint("queue_name", "task_id"),
+    )
+    op.create_index(
+        "ix_tasks_claim",
+        "tasks",
+        ["queue_name", "status", sa.text("priority DESC"), "pending_at_us", "task_id"],
+    )
+    op.create_index("ix_tasks_expiry", "tasks", ["status", "lease_expires_at_us"])
+    op.create_index(
+        "ix_tasks_default_list",
+        "tasks",
+        ["queue_name", sa.text("created_at_us DESC"), sa.text("task_id DESC")],
+    )
+    op.create_index(
+        "ix_tasks_status_list",
+        "tasks",
+        ["queue_name", "status", sa.text("created_at_us DESC"), sa.text("task_id DESC")],
+    )
+    op.create_index(
+        "uq_tasks_active_run_id",
+        "tasks",
+        ["active_run_id"],
+        unique=True,
+        sqlite_where=sa.text("active_run_id IS NOT NULL"),
+    )
+    op.create_index(
+        "ix_tasks_terminal_run_id",
+        "tasks",
+        ["last_terminal_run_id"],
+        sqlite_where=sa.text("last_terminal_run_id IS NOT NULL"),
+    )
+    op.create_table(
+        "task_routes",
+        sa.Column("queue_name", sa.String(length=128), nullable=False),
+        sa.Column("task_id", sa.String(length=14), nullable=False),
+        sa.Column("route", sa.String(length=128), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["queue_name", "task_id"],
+            ["tasks.queue_name", "tasks.task_id"],
+            ondelete="CASCADE",
+        ),
+        sa.PrimaryKeyConstraint("queue_name", "task_id", "route"),
+    )
+    op.create_index(
+        "ix_task_routes_route",
+        "task_routes",
+        ["queue_name", "route", "task_id"],
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("task_routes")
+    op.drop_table("tasks")
+    op.drop_table("queues")

labtasker_server/migrations/versions/__init__.py

@@ -0,0 +1 @@
+"""Labtasker v2 schema revisions."""

labtasker_server/models.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from sqlalchemy import (
+    CheckConstraint,
+    ForeignKeyConstraint,
+    Index,
+    Integer,
+    String,
+    Text,
+    text,
+)
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class QueueRow(Base):
+    __tablename__ = "queues"
+
+    name: Mapped[str] = mapped_column(String(128), primary_key=True)
+
+
+class TaskRow(Base):
+    __tablename__ = "tasks"
+    __table_args__ = (
+        ForeignKeyConstraint(["queue_name"], ["queues.name"], ondelete="CASCADE"),
+        CheckConstraint(
+            "status IN ('pending','running','succeeded','failed','cancelled')",
+            name="ck_tasks_status",
+        ),
+        CheckConstraint("attempt >= 0", name="ck_tasks_attempt_nonnegative"),
+        CheckConstraint("max_attempts > 0", name="ck_tasks_max_attempts_positive"),
+        CheckConstraint(
+            "last_terminal_action IS NULL OR last_terminal_action IN "
+            "('complete','fail','unclaim','heartbeat_expired','cancel')",
+            name="ck_tasks_terminal_action",
+        ),
+        CheckConstraint(
+            "json_valid(args_json) AND json_type(args_json) = 'object'",
+            name="ck_args_json",
+        ),
+        CheckConstraint(
+            "json_valid(metadata_json) AND json_type(metadata_json) = 'object'",
+            name="ck_metadata_json",
+        ),
+        CheckConstraint(
+            "json_valid(result_json) AND json_type(result_json) = 'object'",
+            name="ck_result_json",
+        ),
+        CheckConstraint(
+            "(status = 'pending' AND pending_at_us IS NOT NULL "
+            "AND active_run_id IS NULL AND lease_expires_at_us IS NULL "
+            "AND attempt < max_attempts) OR status != 'pending'",
+            name="ck_tasks_pending_state",
+        ),
+        CheckConstraint(
+            "(status = 'running' AND pending_at_us IS NULL "
+            "AND active_run_id IS NOT NULL AND lease_expires_at_us IS NOT NULL) "
+            "OR status != 'running'",
+            name="ck_tasks_running_state",
+        ),
+        CheckConstraint(
+            "(status IN ('succeeded','failed','cancelled') AND pending_at_us IS NULL "
+            "AND active_run_id IS NULL AND lease_expires_at_us IS NULL) "
+            "OR status NOT IN ('succeeded','failed','cancelled')",
+            name="ck_tasks_terminal_state",
+        ),
+        Index(
+            "ix_tasks_claim",
+            "queue_name",
+            "status",
+            "priority",
+            "pending_at_us",
+            "task_id",
+        ),
+        Index("ix_tasks_expiry", "status", "lease_expires_at_us"),
+        Index("ix_tasks_default_list", "queue_name", "created_at_us", "task_id"),
+        Index("ix_tasks_status_list", "queue_name", "status", "created_at_us", "task_id"),
+        Index(
+            "uq_tasks_active_run_id",
+            "active_run_id",
+            unique=True,
+            sqlite_where=text("active_run_id IS NOT NULL"),
+        ),
+        Index(
+            "ix_tasks_terminal_run_id",
+            "last_terminal_run_id",
+            sqlite_where=text("last_terminal_run_id IS NOT NULL"),
+        ),
+    )
+
+    queue_name: Mapped[str] = mapped_column(String(128), primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(14), primary_key=True)
+    status: Mapped[str] = mapped_column(String(16), nullable=False)
+    name: Mapped[str | None] = mapped_column(String(256))
+    args_json: Mapped[str] = mapped_column(Text, nullable=False)
+    metadata_json: Mapped[str] = mapped_column(Text, nullable=False)
+    result_json: Mapped[str] = mapped_column(Text, nullable=False)
+    priority: Mapped[int] = mapped_column(Integer, nullable=False)
+    attempt: Mapped[int] = mapped_column(Integer, nullable=False)
+    max_attempts: Mapped[int] = mapped_column(Integer, nullable=False)
+    created_at_us: Mapped[int] = mapped_column(Integer, nullable=False)
+    updated_at_us: Mapped[int] = mapped_column(Integer, nullable=False)
+    last_route: Mapped[str | None] = mapped_column(String(128))
+    started_at_us: Mapped[int | None] = mapped_column(Integer)
+    finished_at_us: Mapped[int | None] = mapped_column(Integer)
+    last_error_json: Mapped[str | None] = mapped_column(Text)
+    creation_hash: Mapped[str] = mapped_column(String(64), nullable=False)
+    active_run_id: Mapped[str | None] = mapped_column(String(14))
+    lease_expires_at_us: Mapped[int | None] = mapped_column(Integer)
+    last_terminal_run_id: Mapped[str | None] = mapped_column(String(14))
+    last_terminal_action: Mapped[str | None] = mapped_column(String(32))
+    pending_at_us: Mapped[int | None] = mapped_column(Integer)
+
+    routes: Mapped[list[TaskRouteRow]] = relationship(
+        back_populates="task",
+        cascade="all, delete-orphan",
+        lazy="selectin",
+    )
+
+
+class TaskRouteRow(Base):
+    __tablename__ = "task_routes"
+    __table_args__ = (
+        ForeignKeyConstraint(
+            ["queue_name", "task_id"],
+            ["tasks.queue_name", "tasks.task_id"],
+            ondelete="CASCADE",
+        ),
+        Index("ix_task_routes_route", "queue_name", "route", "task_id"),
+    )
+
+    queue_name: Mapped[str] = mapped_column(String(128), primary_key=True)
+    task_id: Mapped[str] = mapped_column(String(14), primary_key=True)
+    route: Mapped[str] = mapped_column(String(128), primary_key=True)
+
+    task: Mapped[TaskRow] = relationship(back_populates="routes")

labtasker_server/pagination.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import base64
+import binascii
+import json
+from dataclasses import dataclass
+
+from labtasker_server.errors import DomainError, invalid
+from labtasker_server.validation import (
+    INT64_MAX,
+    INT64_MIN,
+    validate_task_id,
+    validate_unicode_scalar,
+)
+
+
+@dataclass(frozen=True, slots=True)
+class TaskSelection:
+    queue: str
+    status: str | None
+    name: str | None
+    filter: str | None
+    order_by: str
+    descending: bool
+
+
+@dataclass(frozen=True, slots=True)
+class CursorPosition:
+    value: str | int | None
+    task_id: str
+
+
+def encode_cursor(selection: TaskSelection, position: CursorPosition) -> str:
+    payload = {
+        "v": 1,
+        "selection": {
+            "queue": selection.queue,
+            "status": selection.status,
+            "name": selection.name,
+            "filter": selection.filter,
+            "order_by": selection.order_by,
+            "descending": selection.descending,
+        },
+        "position": {"value": position.value, "id": position.task_id},
+    }
+    encoded = json.dumps(
+        payload,
+        ensure_ascii=False,
+        allow_nan=False,
+        separators=(",", ":"),
+        sort_keys=True,
+    ).encode("utf-8")
+    return base64.urlsafe_b64encode(encoded).decode("ascii").rstrip("=")
+
+
+def decode_cursor(cursor: str, selection: TaskSelection) -> CursorPosition:
+    try:
+        if not cursor or len(cursor) > 16_384 or not cursor.isascii():
+            raise ValueError
+        padding = "=" * (-len(cursor) % 4)
+        raw = base64.b64decode(cursor + padding, altchars=b"-_", validate=True)
+        payload = json.loads(raw)
+        if not isinstance(payload, dict) or set(payload) != {"v", "selection", "position"}:
+            raise ValueError
+        if payload["v"] != 1 or payload["selection"] != _selection_json(selection):
+            raise ValueError
+        position = payload["position"]
+        if not isinstance(position, dict) or set(position) != {"value", "id"}:
+            raise ValueError
+        task_id = position["id"]
+        value = position["value"]
+        if not isinstance(task_id, str) or isinstance(value, (bool, float, list, dict)):
+            raise ValueError
+        validate_task_id(task_id)
+        _validate_cursor_value(selection.order_by, value)
+        return CursorPosition(value=value, task_id=task_id)
+    except (
+        UnicodeError,
+        ValueError,
+        TypeError,
+        json.JSONDecodeError,
+        binascii.Error,
+        DomainError,
+    ) as error:
+        raise invalid(
+            "invalid_cursor", "Cursor is malformed or does not match this request."
+        ) from error
+
+
+def _selection_json(selection: TaskSelection) -> dict[str, object]:
+    return {
+        "queue": selection.queue,
+        "status": selection.status,
+        "name": selection.name,
+        "filter": selection.filter,
+        "order_by": selection.order_by,
+        "descending": selection.descending,
+    }
+
+
+def _validate_cursor_value(order_by: str, value: object) -> None:
+    numeric_fields = {
+        "priority",
+        "attempt",
+        "max_attempts",
+        "created_at",
+        "updated_at",
+        "started_at",
+        "finished_at",
+    }
+    nullable_fields = {"name", "last_route", "started_at", "finished_at"}
+    if value is None:
+        if order_by not in nullable_fields:
+            raise ValueError
+        return
+    if order_by in numeric_fields:
+        if (
+            not isinstance(value, int)
+            or isinstance(value, bool)
+            or not INT64_MIN <= value <= INT64_MAX
+        ):
+            raise ValueError
+        return
+    if not isinstance(value, str):
+        raise ValueError
+    validate_unicode_scalar(value, field="cursor.position.value")

labtasker_server/py.typed

@@ -0,0 +1 @@
+

labtasker_server/schemas.py

@@ -0,0 +1,265 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from datetime import datetime
+from typing import Annotated, Any, Literal, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field, StrictInt, field_validator, model_validator
+from pydantic_core import PydanticCustomError
+
+from labtasker_server.errors import DomainError
+from labtasker_server.validation import (
+    INT64_MAX,
+    INT64_MIN,
+    JSONValue,
+    canonical_routes,
+    validate_identifier,
+    validate_json_object,
+    validate_run_id,
+    validate_task_name,
+    validate_unicode_scalar,
+)
+
+TaskStatus = Literal["pending", "running", "succeeded", "failed", "cancelled"]
+TaskOrderField = Literal[
+    "id",
+    "name",
+    "status",
+    "priority",
+    "attempt",
+    "max_attempts",
+    "last_route",
+    "created_at",
+    "updated_at",
+    "started_at",
+    "finished_at",
+]
+Int64 = Annotated[StrictInt, Field(ge=INT64_MIN, le=INT64_MAX)]
+PositiveInt64 = Annotated[StrictInt, Field(ge=1, le=INT64_MAX)]
+T = TypeVar("T")
+OMITTED: Any = None
+
+
+class StrictModel(BaseModel):
+    model_config = ConfigDict(extra="forbid", strict=True)
+
+
+class Queue(StrictModel):
+    name: str
+
+
+class LastError(StrictModel):
+    type: str
+    message: str
+    traceback: str | None
+    occurred_at: datetime
+    attempt: int
+    run_id: str
+
+
+class Task(StrictModel):
+    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
+
+
+class TaskCreate(StrictModel):
+    name: str | None = None
+    args: dict[str, JSONValue] = Field(default_factory=dict)
+    metadata: dict[str, JSONValue] = Field(default_factory=dict)
+    priority: Int64 = 0
+    max_attempts: PositiveInt64 = 3
+    routes: list[str] = Field(default_factory=lambda: ["default"])
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, value: str | None) -> str | None:
+        try:
+            return validate_task_name(value)
+        except DomainError as error:
+            raise _pydantic_error(error) from error
+
+    @field_validator("args", "metadata")
+    @classmethod
+    def validate_objects(cls, value: dict[str, JSONValue], info: object) -> dict[str, JSONValue]:
+        field_name = getattr(info, "field_name", "value")
+        try:
+            return validate_json_object(value, field=field_name)
+        except DomainError as error:
+            raise _pydantic_error(error) from error
+
+    @field_validator("routes")
+    @classmethod
+    def validate_routes(cls, value: list[str]) -> list[str]:
+        try:
+            return canonical_routes(value)
+        except DomainError as error:
+            raise _pydantic_error(error) from error
+
+
+class TaskUpdate(StrictModel):
+    name: str | None = None
+    args: dict[str, JSONValue] = OMITTED
+    metadata: dict[str, JSONValue] = OMITTED
+    priority: Int64 = OMITTED
+    max_attempts: PositiveInt64 = OMITTED
+    routes: list[str] = OMITTED
+    result: dict[str, JSONValue] = OMITTED
+
+    @model_validator(mode="after")
+    def validate_nonempty(self) -> TaskUpdate:
+        if not self.model_fields_set:
+            raise PydanticCustomError("invalid_update", "Update must contain at least one field.")
+        return self
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, value: str | None) -> str | None:
+        return _validated(lambda: validate_task_name(value))
+
+    @field_validator("args", "metadata", "result")
+    @classmethod
+    def validate_objects(
+        cls,
+        value: dict[str, JSONValue],
+        info: object,
+    ) -> dict[str, JSONValue]:
+        field_name = getattr(info, "field_name", "value")
+        return _validated(lambda: validate_json_object(value, field=field_name))
+
+    @field_validator("routes")
+    @classmethod
+    def validate_routes(cls, value: list[str]) -> list[str]:
+        return _validated(lambda: canonical_routes(value))
+
+
+class BulkUpdateRequest(StrictModel):
+    filter: str
+    changes: TaskUpdate
+
+
+class BulkUpdateResult(StrictModel):
+    matched: int
+    updated: int
+
+
+class TaskPage(StrictModel):
+    items: list[Task]
+    next_cursor: str | None
+
+
+class CountResponse(StrictModel):
+    count: int
+
+
+class ClaimRequest(StrictModel):
+    route: str
+    run_id: str
+
+    @field_validator("route")
+    @classmethod
+    def validate_route(cls, value: str) -> str:
+        return _validated(lambda: validate_identifier(value, kind="Route"))
+
+    @field_validator("run_id")
+    @classmethod
+    def validate_run(cls, value: str) -> str:
+        return _validated(lambda: validate_run_id(value))
+
+
+class ClaimResponse(StrictModel):
+    task: Task
+    run_id: str
+    lease_expires_at: datetime
+
+
+class RunRequest(StrictModel):
+    run_id: str
+
+    @field_validator("run_id")
+    @classmethod
+    def validate_run(cls, value: str) -> str:
+        return _validated(lambda: validate_run_id(value))
+
+
+class HeartbeatResponse(StrictModel):
+    lease_expires_at: datetime
+
+
+class CompleteRequest(RunRequest):
+    result: dict[str, JSONValue]
+
+    @field_validator("result")
+    @classmethod
+    def validate_result(cls, value: dict[str, JSONValue]) -> dict[str, JSONValue]:
+        return _validated(lambda: validate_json_object(value, field="result"))
+
+
+class FailureReport(StrictModel):
+    type: str
+    message: str
+    traceback: str | None
+
+    @field_validator("type", "message", "traceback")
+    @classmethod
+    def validate_strings(cls, value: str | None, info: object) -> str | None:
+        if value is None:
+            return None
+        field_name = getattr(info, "field_name", "error")
+        return _validated(lambda: validate_unicode_scalar(value, field=field_name))
+
+
+class FailRequest(RunRequest):
+    error: FailureReport
+
+
+class ErrorItem(StrictModel):
+    location: list[str | int]
+    message: str
+
+
+class ErrorBody(StrictModel):
+    code: str
+    message: str
+    details: dict[str, JSONValue]
+
+
+class ErrorEnvelope(StrictModel):
+    error: ErrorBody
+
+
+class HealthyResponse(StrictModel):
+    status: Literal["ok"]
+    api_version: Literal["2"]
+    database: Literal["ok"]
+
+
+class UnhealthyResponse(StrictModel):
+    status: Literal["error"]
+    api_version: Literal["2"]
+    database: Literal["error"]
+
+
+def _pydantic_error(error: DomainError) -> PydanticCustomError:
+    return PydanticCustomError(error.code, error.message, error.details)
+
+
+def _validated(operation: Callable[[], T]) -> T:
+    try:
+        return operation()
+    except DomainError as error:
+        raise _pydantic_error(error) from error

labtasker_server/services/__init__.py

@@ -0,0 +1 @@
+"""Application services owning Labtasker transactions and domain behavior."""