Merge pull request #60 from taskbadger/sk/parents

Support parent/child tasks

Author: snopoke

integration_tests/tasks.py

@@ -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."""

integration_tests/test_celery.py

@@ -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)
@@ -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)

integration_tests/test_parents.py

@@ -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)

pyproject.toml

@@ -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]",
 ]

taskbadger.yaml

@@ -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:
@@ -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
@@ -420,6 +428,7 @@ paths:
       security:
       - projectKeyAuth: []
       - bearerAuth: []
+      deprecated: true
       responses:
         '201':
           content:
@@ -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
@@ -509,6 +521,7 @@ paths:
       security:
       - projectKeyAuth: []
       - bearerAuth: []
+      deprecated: true
       responses:
         '200':
           content:
@@ -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
@@ -555,6 +571,7 @@ paths:
       security:
       - projectKeyAuth: []
       - bearerAuth: []
+      deprecated: true
       responses:
         '200':
           content:
@@ -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
@@ -596,6 +616,7 @@ paths:
       security:
       - projectKeyAuth: []
       - bearerAuth: []
+      deprecated: true
       responses:
         '204':
           description: No response body
@@ -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
@@ -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
@@ -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

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, get_task, init, list_tasks, update_task
 
 __all__ = [
     "track",
@@ -20,6 +20,7 @@
     "create_task",
     "get_task",
     "init",
+    "list_tasks",
     "update_task",
 ]
 

taskbadger/_current_task.py

@@ -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

taskbadger/_integrations.py

@@ -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.