Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions integration_tests/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,33 @@ def add_auto_track(self, x, y):
return x + y


@shared_task(bind=True, base=taskbadger.celery.Task)
def spawns_grandchild(self, x, y):
"""Deferred by `spawns_child`, and defers a task of its own in turn.

Since this is already a child, the task it defers has to be flattened onto
the root rather than nested under this one — the API only allows one level.
"""
grandchild = add.delay(x, y)
return {"own_tb_id": self.taskbadger_task_id, "grandchild_tb_id": grandchild.taskbadger_task_id}


@shared_task(bind=True, base=taskbadger.celery.Task)
def spawns_child(self, x, y):
"""Defers a task from inside its own run, so that task nests under this one."""
return spawns_grandchild.delay(x, y).id


@shared_task(bind=True, base=taskbadger.celery.Task)
def chain_head(self):
return self.taskbadger_task_id


@shared_task(bind=True, base=taskbadger.celery.Task)
def chain_tail(self, head_tb_id):
return {"head_tb_id": head_tb_id, "own_tb_id": self.taskbadger_task_id}


@shared_task(bind=True, base=taskbadger.celery.Task, taskbadger_heartbeat_interval=HEARTBEAT_INTERVAL)
def slow_add(self, x, y):
"""Runs long enough to go stale without a heartbeat, and never updates itself."""
Expand Down
39 changes: 38 additions & 1 deletion integration_tests/test_celery.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,12 @@
import time

import pytest
from celery import chain

import taskbadger
from taskbadger import StatusEnum

from .tasks import HEARTBEAT_INTERVAL, add, add_auto_track, slow_add
from .tasks import HEARTBEAT_INTERVAL, add, add_auto_track, chain_head, chain_tail, slow_add, spawns_child


@pytest.fixture(autouse=True)
Expand Down Expand Up @@ -45,6 +46,42 @@ def test_celery_auto_track(celery_session_app, celery_session_worker):
assert result.get(timeout=10, propagate=True) == a + b


def test_celery_child_task_nests_under_its_parent(celery_session_app, celery_session_worker):
a, b = random.randint(1, 1000), random.randint(1, 1000)
root = spawns_child.delay(a, b)
child_celery_id = root.get(timeout=15, propagate=True)

child = celery_session_app.AsyncResult(child_celery_id).get(timeout=15, propagate=True)

assert taskbadger.get_task(child["own_tb_id"]).parent == root.taskbadger_task_id


def test_celery_grandchild_is_flattened_onto_the_root(celery_session_app, celery_session_worker):
"""Nesting stops at one level, so a task deferred by a child joins it under
the root instead of hanging off it (which the API would reject)."""
a, b = random.randint(1, 1000), random.randint(1, 1000)
root = spawns_child.delay(a, b)
child_celery_id = root.get(timeout=15, propagate=True)

child = celery_session_app.AsyncResult(child_celery_id).get(timeout=15, propagate=True)
grandchild = taskbadger.get_task(child["grandchild_tb_id"])

assert grandchild.parent == root.taskbadger_task_id
assert grandchild.parent != child["own_tb_id"]


def test_celery_chain_links_are_not_nested(celery_session_app, celery_session_worker):
"""Celery dispatches the next chain link from inside the previous task's run,
so it would otherwise be nested under it. Links are successors, not subtasks.
"""
ids = chain(chain_head.s(), chain_tail.s()).apply_async().get(timeout=20, propagate=True)

assert ids["head_tb_id"], "the first link should be tracked"
assert ids["own_tb_id"], "the second link should be tracked"
assert not taskbadger.get_task(ids["own_tb_id"]).parent
assert taskbadger.list_tasks(parent=ids["head_tb_id"]).results == []


def test_celery_heartbeat(celery_session_app, celery_session_worker):
"""The worker pings the task while it runs, so it doesn't go stale."""
a, b = random.randint(1, 1000), random.randint(1, 1000)
Expand Down
45 changes: 45 additions & 0 deletions integration_tests/test_parents.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
import pytest

import taskbadger as badger
from taskbadger.exceptions import UnexpectedStatus


def test_create_child_task():
parent = badger.create_task("test parent")
child = badger.create_task("test child", parent=parent.id)

assert child.parent == parent.id
assert badger.get_task(child.id).parent == parent.id
# the parent itself stays a root
assert not badger.get_task(parent.id).parent


def test_list_tasks_by_parent():
parent = badger.create_task("test parent for listing")
child = badger.create_task("test child for listing", parent=parent.id)
badger.create_task("test unrelated task")

children = badger.list_tasks(parent=parent.id).results

assert [task.id for task in children] == [child.id]


def test_set_parent_on_an_existing_task():
parent = badger.create_task("test parent for update")
child = badger.create_task("test child for update")
assert not child.parent

child.update(parent=parent.id)

assert child.parent == parent.id
assert badger.get_task(child.id).parent == parent.id


def test_nesting_is_limited_to_one_level():
"""The API rejects a grandchild, which is what the integrations' flattening
exists to avoid."""
parent = badger.create_task("test parent depth")
child = badger.create_task("test child depth", parent=parent.id)

with pytest.raises(UnexpectedStatus):
badger.create_task("test grandchild depth", parent=child.id)
4 changes: 3 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,9 @@ dev = [
"invoke",
"pytest-celery",
"redis",
"openapi-python-client",
# 0.29 generates `datetime.fromisoformat` calls, which can't parse the API's
# `Z`-suffixed timestamps on Python 3.10.
"openapi-python-client<0.29",
"taskbadger[cli]",
"taskbadger[sentry]",
]
Expand Down
46 changes: 42 additions & 4 deletions taskbadger.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,11 @@ paths:
description: Number of results to return per page.
schema:
type: integer
- in: query
name: parent
schema:
type: string
description: Only return the tasks that are part of this task.
- in: path
name: project_slug
schema:
Expand Down Expand Up @@ -388,7 +393,10 @@ paths:
description: ''
post:
operationId: action_create
description: Create an action for a task
description: '**Deprecated.** Per-job actions are being retired. Configure alerts
once at the project level with a global action instead of attaching them to
individual tasks. This endpoint still works but will be removed in a future
release.'
summary: Create Action
parameters:
- in: path
Expand Down Expand Up @@ -420,6 +428,7 @@ paths:
security:
- projectKeyAuth: []
- bearerAuth: []
deprecated: true
responses:
'201':
content:
Expand Down Expand Up @@ -471,7 +480,10 @@ paths:
description: ''
put:
operationId: action_update
description: Update an action
description: '**Deprecated.** Per-job actions are being retired. Configure alerts
once at the project level with a global action instead of attaching them to
individual tasks. This endpoint still works but will be removed in a future
release.'
summary: Update Action
parameters:
- in: path
Expand Down Expand Up @@ -509,6 +521,7 @@ paths:
security:
- projectKeyAuth: []
- bearerAuth: []
deprecated: true
responses:
'200':
content:
Expand All @@ -518,7 +531,10 @@ paths:
description: ''
patch:
operationId: action_partial_update
description: Update an action
description: '**Deprecated.** Per-job actions are being retired. Configure alerts
once at the project level with a global action instead of attaching them to
individual tasks. This endpoint still works but will be removed in a future
release.'
summary: Update Action (partial)
parameters:
- in: path
Expand Down Expand Up @@ -555,6 +571,7 @@ paths:
security:
- projectKeyAuth: []
- bearerAuth: []
deprecated: true
responses:
'200':
content:
Expand All @@ -564,7 +581,10 @@ paths:
description: ''
delete:
operationId: action_cancel
description: Cancel an action
description: '**Deprecated.** Per-job actions are being retired. Configure alerts
once at the project level with a global action instead of attaching them to
individual tasks. This endpoint still works but will be removed in a future
release.'
summary: Cancel Action
parameters:
- in: path
Expand Down Expand Up @@ -596,6 +616,7 @@ paths:
security:
- projectKeyAuth: []
- bearerAuth: []
deprecated: true
responses:
'204':
description: No response body
Expand Down Expand Up @@ -685,6 +706,12 @@ components:
PatchedTaskRequest:
type: object
properties:
parent:
type: string
minLength: 1
nullable: true
description: ID of the task this task is part of. Tasks can only be nested
one level deep, and a task's parent can not be changed once set.
name:
type: string
minLength: 1
Expand Down Expand Up @@ -791,6 +818,11 @@ components:
project:
type: string
readOnly: true
parent:
type: string
nullable: true
description: ID of the task this task is part of. Tasks can only be nested
one level deep, and a task's parent can not be changed once set.
name:
type: string
description: Name of the task
Expand Down Expand Up @@ -894,6 +926,12 @@ components:
TaskRequest:
type: object
properties:
parent:
type: string
minLength: 1
nullable: true
description: ID of the task this task is part of. Tasks can only be nested
one level deep, and a task's parent can not be changed once set.
name:
type: string
minLength: 1
Expand Down
3 changes: 2 additions & 1 deletion taskbadger/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, get_task, init, list_tasks, update_task

__all__ = [
"track",
Expand All @@ -20,6 +20,7 @@
"create_task",
"get_task",
"init",
"list_tasks",
"update_task",
]

Expand Down
53 changes: 53 additions & 0 deletions taskbadger/_current_task.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
"""Tracks which tracked task is running in the current context so that tasks
created while it runs can be attached to it as children. Not part of the public
API.

Task Badger nests tasks a single level deep, so `parent_id` never returns the
id of a task that is itself a child: entering a task that already has a parent
keeps that parent as the id offered for nesting. A task deferred by a child
therefore lands alongside it under the same root rather than being rejected.

Call sites resolve a task's own parent themselves — they all have it to hand
already (from the create response, or from the cache the status update fills),
which keeps this module free of API calls.
"""

from contextvars import ContextVar

# (id of the running task, id of its parent or None if it is a root task)
_current: ContextVar[tuple[str, str | None] | None] = ContextVar("taskbadger_current_task", default=None)


def enter_task(task_id: str, parent: str = None):
"""Mark `task_id` as the task running in this context.

Arguments:
task_id: The running task.
parent: The running task's own parent, if it has one.

Returns:
A token to pass to `exit_task`.
"""
return _current.set((task_id, parent))


def exit_task(token) -> None:
_current.reset(token)


def current_task_id() -> str | None:
"""The id of the tracked task running in this context, if any."""
current = _current.get()
return current[0] if current else None


def parent_id() -> str | None:
"""The id a task created right now should use as its `parent`.

`None` outside a tracked task.
"""
current = _current.get()
if current is None:
return None
task_id, parent = current
return parent or task_id
11 changes: 11 additions & 0 deletions taskbadger/_integrations.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,17 @@ def safe_get_task(task_id: str):
return task


def parent_of(task) -> str | None:
"""The id of ``task``'s parent, or ``None`` if it has none.

Accepts ``None`` (e.g. a failed fetch) and normalizes the generated model's
``UNSET`` — returned for tasks fetched before the API grew the field — to
``None``.
"""
parent = getattr(task, "parent", None) if task is not None else None
return parent or None


def match_task_name(task_name: str, includes, excludes) -> bool:
"""Return True if ``task_name`` should be tracked under the given rules.

Expand Down
Loading