diff --git a/taskbadger.yaml b/taskbadger.yaml index 750a24b..1e29726 100644 --- a/taskbadger.yaml +++ b/taskbadger.yaml @@ -685,6 +685,11 @@ components: PatchedTaskRequest: type: object properties: + id: + type: string + minLength: 1 + description: Task ID. May be set on creation to a UUID or shortened UUID; + it must be unique and is immutable thereafter. If omitted, an ID is generated. name: type: string minLength: 1 @@ -783,8 +788,8 @@ components: properties: id: type: string - readOnly: true - description: Task ID + description: Task ID. May be set on creation to a UUID or shortened UUID; + it must be unique and is immutable thereafter. If omitted, an ID is generated. organization: type: string readOnly: true @@ -883,7 +888,6 @@ components: to 'value'. required: - created - - id - name - organization - project @@ -894,6 +898,11 @@ components: TaskRequest: type: object properties: + id: + type: string + minLength: 1 + description: Task ID. May be set on creation to a UUID or shortened UUID; + it must be unique and is immutable thereafter. If omitted, an ID is generated. name: type: string minLength: 1 diff --git a/taskbadger/__init__.py b/taskbadger/__init__.py index 9f541f3..84feba5 100644 --- a/taskbadger/__init__.py +++ b/taskbadger/__init__.py @@ -3,7 +3,7 @@ from .internal.models import StatusEnum from .mug import Badger, Session from .safe_sdk import create_task_safe, update_task_safe -from .sdk import DefaultMergeStrategy, Task, create_task, get_task, init, update_task +from .sdk import DefaultMergeStrategy, Task, create_task, generate_task_id, get_task, init, update_task __all__ = [ "track", @@ -18,6 +18,7 @@ "DefaultMergeStrategy", "Task", "create_task", + "generate_task_id", "get_task", "init", "update_task", diff --git a/taskbadger/internal/models/patched_task_request.py b/taskbadger/internal/models/patched_task_request.py index 40a7d72..5cd61a6 100644 --- a/taskbadger/internal/models/patched_task_request.py +++ b/taskbadger/internal/models/patched_task_request.py @@ -22,6 +22,8 @@ class PatchedTaskRequest: """ Attributes: + id (str | Unset): Task ID. May be set on creation to a UUID or shortened UUID; it must be unique and is + immutable thereafter. If omitted, an ID is generated. name (str | Unset): Name of the task queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with @@ -50,6 +52,7 @@ class PatchedTaskRequest: tags (PatchedTaskRequestTags | Unset): Tags for the task represented as a mapping from 'namespace' to 'value'. """ + id: str | Unset = UNSET name: str | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET @@ -68,6 +71,8 @@ class PatchedTaskRequest: def to_dict(self) -> dict[str, Any]: from ..models.patched_task_request_tags import PatchedTaskRequestTags + id = self.id + name = self.name queue = self.queue @@ -129,6 +134,8 @@ def to_dict(self) -> dict[str, Any]: field_dict: dict[str, Any] = {} field_dict.update(self.additional_properties) field_dict.update({}) + if id is not UNSET: + field_dict["id"] = id if name is not UNSET: field_dict["name"] = name if queue is not UNSET: @@ -163,6 +170,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T: from ..models.patched_task_request_tags import PatchedTaskRequestTags d = dict(src_dict) + id = d.pop("id", UNSET) + name = d.pop("name", UNSET) queue = d.pop("queue", UNSET) @@ -258,6 +267,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: tags = PatchedTaskRequestTags.from_dict(_tags) patched_task_request = cls( + id=id, name=name, queue=queue, external_id=external_id, diff --git a/taskbadger/internal/models/task.py b/taskbadger/internal/models/task.py index 718616b..e392e00 100644 --- a/taskbadger/internal/models/task.py +++ b/taskbadger/internal/models/task.py @@ -22,7 +22,6 @@ class Task: """ Attributes: - id (str): Task ID organization (str): project (str): name (str): Name of the task @@ -31,6 +30,8 @@ class Task: updated (datetime.datetime): url (str): public_url (str): + id (str | Unset): Task ID. May be set on creation to a UUID or shortened UUID; it must be unique and is + immutable thereafter. If omitted, an ID is generated. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -58,7 +59,6 @@ class Task: tags (TaskTags | Unset): Tags for the task represented as a mapping from 'namespace' to 'value'. """ - id: str organization: str project: str name: str @@ -67,6 +67,7 @@ class Task: updated: datetime.datetime url: str public_url: str + id: str | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING @@ -84,8 +85,6 @@ class Task: def to_dict(self) -> dict[str, Any]: from ..models.task_tags import TaskTags - id = self.id - organization = self.organization project = self.project @@ -103,6 +102,8 @@ def to_dict(self) -> dict[str, Any]: public_url = self.public_url + id = self.id + queue = self.queue external_id = self.external_id @@ -163,7 +164,6 @@ def to_dict(self) -> dict[str, Any]: field_dict.update(self.additional_properties) field_dict.update( { - "id": id, "organization": organization, "project": project, "name": name, @@ -174,6 +174,8 @@ def to_dict(self) -> dict[str, Any]: "public_url": public_url, } ) + if id is not UNSET: + field_dict["id"] = id if queue is not UNSET: field_dict["queue"] = queue if external_id is not UNSET: @@ -206,8 +208,6 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T: from ..models.task_tags import TaskTags d = dict(src_dict) - id = d.pop("id") - organization = d.pop("organization") project = d.pop("project") @@ -229,6 +229,8 @@ def _parse_value_percent(data: object) -> int | None: public_url = d.pop("public_url") + id = d.pop("id", UNSET) + queue = d.pop("queue", UNSET) external_id = d.pop("external_id", UNSET) @@ -322,7 +324,6 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: tags = TaskTags.from_dict(_tags) task = cls( - id=id, organization=organization, project=project, name=name, @@ -331,6 +332,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: updated=updated, url=url, public_url=public_url, + id=id, queue=queue, external_id=external_id, status=status, diff --git a/taskbadger/internal/models/task_request.py b/taskbadger/internal/models/task_request.py index 363630c..3031fbe 100644 --- a/taskbadger/internal/models/task_request.py +++ b/taskbadger/internal/models/task_request.py @@ -23,6 +23,8 @@ class TaskRequest: """ Attributes: name (str): Name of the task + id (str | Unset): Task ID. May be set on creation to a UUID or shortened UUID; it must be unique and is + immutable thereafter. If omitted, an ID is generated. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -51,6 +53,7 @@ class TaskRequest: """ name: str + id: str | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING @@ -70,6 +73,8 @@ def to_dict(self) -> dict[str, Any]: name = self.name + id = self.id + queue = self.queue external_id = self.external_id @@ -133,6 +138,8 @@ def to_dict(self) -> dict[str, Any]: "name": name, } ) + if id is not UNSET: + field_dict["id"] = id if queue is not UNSET: field_dict["queue"] = queue if external_id is not UNSET: @@ -167,6 +174,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T: d = dict(src_dict) name = d.pop("name") + id = d.pop("id", UNSET) + queue = d.pop("queue", UNSET) external_id = d.pop("external_id", UNSET) @@ -261,6 +270,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: task_request = cls( name=name, + id=id, queue=queue, external_id=external_id, status=status, diff --git a/taskbadger/procrastinate.py b/taskbadger/procrastinate.py index 3ae6364..156c9c6 100644 --- a/taskbadger/procrastinate.py +++ b/taskbadger/procrastinate.py @@ -21,7 +21,7 @@ from .internal.models import StatusEnum from .mug import Badger from .safe_sdk import create_task_safe, update_task_safe -from .sdk import DefaultMergeStrategy +from .sdk import DefaultMergeStrategy, generate_task_id log = logging.getLogger("taskbadger") @@ -164,7 +164,7 @@ async def defer_async(**kwargs): task.defer_async = defer_async -def _create_pending_task(task, task_kwargs, queue=None): +def _create_pending_task(task, task_kwargs, queue=None, task_id=None): """Create a PENDING TaskBadger task for ``task`` if it should be tracked. Returns the created TaskBadger task, or ``None`` if Badger isn't @@ -172,6 +172,7 @@ def _create_pending_task(task, task_kwargs, queue=None): create call failed. ``task_kwargs`` is used only for the ``record_task_args`` data capture. ``queue`` overrides the queue name recorded on the TaskBadger task (defaults to the task's own queue). + ``task_id`` assigns the task's id (see ``_maybe_create_pending``). """ if not Badger.is_configured(): return None @@ -203,18 +204,22 @@ def _create_pending_task(task, task_kwargs, queue=None): if data: create_kwargs["data"] = data - return create_task_safe(name, **create_kwargs) + return create_task_safe(name, task_id=task_id, **create_kwargs) def _maybe_create_pending(task, kwargs): """Decide whether to track this defer, and if so create the TaskBadger - task and inject its id into ``kwargs``. Always returns the kwargs dict.""" - tb_task = _create_pending_task(task, kwargs) + task and inject its id into ``kwargs``. Always returns the kwargs dict. + + The id is generated up front so it's known before the task exists, letting us + inject it without depending on the create response.""" + tb_id = generate_task_id() + tb_task = _create_pending_task(task, kwargs, task_id=tb_id) if tb_task is None: return kwargs new_kwargs = dict(kwargs) - new_kwargs[TB_TASK_ID_KWARG] = tb_task.id + new_kwargs[TB_TASK_ID_KWARG] = tb_id return new_kwargs @@ -314,11 +319,12 @@ async def patched(*, job, periodic_id, defer_timestamp): task = app.tasks.get(job.task_name) tb_id = None if task is not None: - tb_task = _create_pending_task(task, job.task_kwargs, queue=job.queue) + candidate_id = generate_task_id() + tb_task = _create_pending_task(task, job.task_kwargs, queue=job.queue, task_id=candidate_id) if tb_task is not None: - new_kwargs = {**job.task_kwargs, TB_TASK_ID_KWARG: tb_task.id} + new_kwargs = {**job.task_kwargs, TB_TASK_ID_KWARG: candidate_id} job = job.evolve(task_kwargs=new_kwargs) - tb_id = tb_task.id + tb_id = candidate_id job_id = await jm._taskbadger_original_defer_periodic_job( job=job, periodic_id=periodic_id, defer_timestamp=defer_timestamp ) diff --git a/taskbadger/sdk.py b/taskbadger/sdk.py index 4389cf2..a275ac4 100644 --- a/taskbadger/sdk.py +++ b/taskbadger/sdk.py @@ -2,6 +2,7 @@ import datetime import logging import os +import uuid import warnings from typing import Any @@ -140,6 +141,33 @@ def get_task(task_id: str) -> "Task": return Task(task) +# Flickr Base58 alphabet, mirroring the server's short-ID encoding. +_SHORT_ID_ALPHABET = "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ" + + +def generate_task_id() -> str: + """Generate a task ID. + + Returns an ID in the same canonical short form the server assigns. Pass it to + ``create_task(task_id=...)`` when you need the task's ID before creating the task, + e.g. to correlate a queued job with its task. + """ + return _shorten_uuid(uuid.uuid4()) + + +def _shorten_uuid(value: uuid.UUID) -> str: + """Encode a UUID into the server's canonical short-ID form: a 4-character hex prefix + followed by the Flickr Base58 encoding of the UUID.""" + base = len(_SHORT_ID_ALPHABET) + number = value.int + chars = [] + while number > 0: + number, remainder = divmod(number, base) + chars.append(_SHORT_ID_ALPHABET[remainder]) + body = "".join(reversed(chars)) or _SHORT_ID_ALPHABET[0] + return f"{value.hex[:4]}{body}" + + def create_task( name: str, status: StatusEnum = StatusEnum.PENDING, @@ -152,6 +180,7 @@ def create_task( monitor_id: str = None, tags: dict[str, str] = None, queue: str = None, + task_id: str = None, external_id: str = None, ) -> "Task": """Create a Task. @@ -168,6 +197,8 @@ def create_task( monitor_id: ID of the monitor to associate this task with. tags: Dictionary of namespace -> value tags. queue: Name of the queue the task is from. + task_id: ID to assign to the task. May be a UUID or shortened UUID; it must be + unique and is immutable. If omitted, an ID is generated by the server. external_id: Identifier from the originating system (e.g. Celery task ID) for correlating with logs. Returns: @@ -177,6 +208,8 @@ def create_task( "name": name, "status": status, } + if task_id is not None: + task_dict["id"] = task_id if queue is not None: task_dict["queue"] = queue if external_id is not None: @@ -343,6 +376,7 @@ def create( monitor_id: str = None, tags: dict[str, str] = None, queue: str = None, + task_id: str = None, external_id: str = None, ) -> "Task": """Create a new task @@ -361,6 +395,7 @@ def create( monitor_id=monitor_id, tags=tags, queue=queue, + task_id=task_id, external_id=external_id, ) diff --git a/tests/test_procrastinate.py b/tests/test_procrastinate.py index 7854261..ed4ffd0 100644 --- a/tests/test_procrastinate.py +++ b/tests/test_procrastinate.py @@ -123,12 +123,13 @@ def add3(a, b): create.assert_called_once() assert create.call_args.args == ("add3",) - assert create.call_args.kwargs == {"status": StatusEnum.PENDING, "queue": "default"} + task_id = create.call_args.kwargs["task_id"] + assert create.call_args.kwargs == {"status": StatusEnum.PENDING, "queue": "default", "task_id": task_id} - # The injected id should appear in the Procrastinate job's task kwargs. + # The generated id is passed to create and injected into the Procrastinate job's task kwargs. jobs = list(app.connector.jobs.values()) assert len(jobs) == 1 - assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id + assert jobs[0]["args"][TB_TASK_ID_KWARG] == task_id @pytest.mark.usefixtures("_bind_settings") @@ -175,13 +176,13 @@ async def add5(a, b): tb = task_for_test() with ( - mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb), + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create, mock.patch("taskbadger.procrastinate.update_task_safe"), ): asyncio.run(add5.defer_async(a=1, b=2)) jobs = list(app.connector.jobs.values()) - assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id + assert jobs[0]["args"][TB_TASK_ID_KWARG] == create.call_args.kwargs["task_id"] @pytest.mark.usefixtures("_bind_settings") @@ -194,12 +195,12 @@ def add_ext(a, b): tb = task_for_test() with ( - mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb), + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create, mock.patch("taskbadger.procrastinate.update_task_safe") as update, ): job_id = add_ext.defer(a=1, b=2) - update.assert_called_once_with(tb.id, external_id=str(job_id)) + update.assert_called_once_with(create.call_args.kwargs["task_id"], external_id=str(job_id)) @pytest.mark.usefixtures("_bind_settings") @@ -212,12 +213,12 @@ async def add_ext_async(a, b): tb = task_for_test() with ( - mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb), + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create, mock.patch("taskbadger.procrastinate.update_task_safe") as update, ): job_id = asyncio.run(add_ext_async.defer_async(a=1, b=2)) - update.assert_called_once_with(tb.id, external_id=str(job_id)) + update.assert_called_once_with(create.call_args.kwargs["task_id"], external_id=str(job_id)) def test_defer_no_external_id_when_untracked(app): @@ -269,7 +270,7 @@ def bare(a): tb = task_for_test() with ( - mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb), + mock.patch("taskbadger.procrastinate.create_task_safe", return_value=tb) as create, mock.patch("taskbadger.procrastinate.update_task_safe"), ): bare.defer(a=1) @@ -277,7 +278,7 @@ def bare(a): assert getattr(bare, "_taskbadger_manual") is True # Inspect the actual Procrastinate job - jobs is a dict keyed by int, kwargs under "args" jobs = list(app.connector.jobs.values()) - assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id + assert jobs[0]["args"][TB_TASK_ID_KWARG] == create.call_args.kwargs["task_id"] @pytest.mark.usefixtures("_bind_settings") @@ -296,12 +297,14 @@ def raw(a): create.assert_called_once() assert create.call_args.args == ("custom",) + task_id = create.call_args.kwargs["task_id"] assert create.call_args.kwargs == { "status": StatusEnum.PENDING, "value_max": 10, "tags": {"env": "test"}, "data": {"k": "v"}, "queue": "default", + "task_id": task_id, } diff --git a/tests/test_procrastinate_system_integration.py b/tests/test_procrastinate_system_integration.py index 057053d..40ed414 100644 --- a/tests/test_procrastinate_system_integration.py +++ b/tests/test_procrastinate_system_integration.py @@ -65,7 +65,7 @@ def auto_target(a): create.assert_called_once() # InMemoryConnector.jobs is a dict keyed by int; kwargs under "args" jobs = list(app.connector.jobs.values()) - assert jobs[0]["args"][TB_TASK_ID_KWARG] == tb.id + assert jobs[0]["args"][TB_TASK_ID_KWARG] == create.call_args.kwargs["task_id"] @pytest.mark.usefixtures("_bind_settings") @@ -121,7 +121,7 @@ def periodic_target(timestamp): create.assert_called_once() jobs_stored = list(app.connector.jobs.values()) - assert jobs_stored[0]["args"][TB_TASK_ID_KWARG] == tb.id + assert jobs_stored[0]["args"][TB_TASK_ID_KWARG] == create.call_args.kwargs["task_id"] @pytest.mark.usefixtures("_bind_settings") diff --git a/tests/test_sdk.py b/tests/test_sdk.py index 768cade..99245e3 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -1,4 +1,5 @@ import datetime +import uuid import warnings from http import HTTPStatus from unittest import mock @@ -13,7 +14,7 @@ ) from taskbadger.internal.types import UNSET, Response from taskbadger.mug import Badger -from taskbadger.sdk import Task, init +from taskbadger.sdk import Task, _shorten_uuid, generate_task_id, init from tests.utils import task_for_test @@ -114,6 +115,33 @@ def test_create_with_queue(settings, patched_create): ) +def test_create_with_task_id(settings, patched_create): + api_task = task_for_test() + patched_create.return_value = Response(HTTPStatus.OK, b"", {}, api_task) + + task_id = "b3f8c1e2-1234-4a5b-8c9d-0e1f2a3b4c5d" + Task.create(name="task name", task_id=task_id) + + request = TaskRequest(name="task name", status=StatusEnum.PENDING, id=task_id) + patched_create.assert_called_with( + client=mock.ANY, + organization_slug="org", + project_slug="project", + body=request, + ) + + +def test_generate_task_id(): + # Matches the server's shorten_uuid: 4-char hex prefix + Flickr Base58 body. + assert _shorten_uuid(uuid.UUID("b3f8c1e2-1234-4a5b-8c9d-0e1f2a3b4c5d")) == "b3f8odYrrReCZKKtXeK82DnyfZ" + + task_id = generate_task_id() + assert task_id[:4] == task_id[:4].lower() + assert all(c in "0123456789abcdef" for c in task_id[:4]) + assert all(c in "123456789abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ" for c in task_id[4:]) + assert generate_task_id() != generate_task_id() + + def test_update_queue(settings, patched_update): api_task = task_for_test() task = Task(api_task) diff --git a/tests/utils.py b/tests/utils.py index 34d9d65..58f8d8e 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -15,9 +15,9 @@ def task_for_test(**kwargs): kwargs["created"] = datetime.datetime.now(datetime.timezone.utc) kwargs["updated"] = datetime.datetime.now(datetime.timezone.utc) return TaskInternal( - task_id, "org", "project", "task_name", + id=task_id, **kwargs, ) diff --git a/uv.lock b/uv.lock index f791468..4394d52 100644 --- a/uv.lock +++ b/uv.lock @@ -1176,7 +1176,7 @@ wheels = [ [[package]] name = "taskbadger" -version = "2.1.0a2" +version = "2.2.0" source = { editable = "." } dependencies = [ { name = "attrs" },