feat: add helper to walk up a task graph through the Taskcluster API (#…

…418) This can be useful in places like action tasks where you don't have access to a generated task graph.
taskcluster · Jan 31, 2024 · c4055e0 · c4055e0
1 parent f277b67
commit c4055e0
Show file tree

Hide file tree

Showing 2 changed files with 184 additions and 0 deletions.
diff --git a/src/taskgraph/util/taskcluster.py b/src/taskgraph/util/taskcluster.py
@@ -7,6 +7,7 @@
 import functools
 import logging
 import os
+from typing import Dict, List, Union
 
 import requests
 import taskcluster_urls as liburls
@@ -367,3 +368,39 @@ def list_task_group_incomplete_tasks(task_group_id):
             params = {"continuationToken": resp.get("continuationToken")}
         else:
             break
+
+
+def _get_deps(task_ids, use_proxy):
+    upstream_tasks = {}
+    for task_id in task_ids:
+        task_def = get_task_definition(task_id, use_proxy)
+        upstream_tasks[task_def["metadata"]["name"]] = task_id
+
+        upstream_tasks.update(_get_deps(task_def["dependencies"], use_proxy))
+
+    return upstream_tasks
+
+
+def get_ancestors(
+    task_ids: Union[List[str], str], use_proxy: bool = False
+) -> Dict[str, str]:
+    """Gets the ancestor tasks of the given task_ids as a dictionary of label -> taskid.
+
+    Args:
+        task_ids (str or [str]): A single task id or a list of task ids to find the ancestors of.
+        use_proxy (bool): See get_root_url.
+
+    Returns:
+        dict: A dict whose keys are task labels and values are task ids.
+    """
+    upstream_tasks: Dict[str, str] = {}
+
+    if isinstance(task_ids, str):
+        task_ids = [task_ids]
+
+    for task_id in task_ids:
+        task_def = get_task_definition(task_id, use_proxy)
+
+        upstream_tasks.update(_get_deps(task_def["dependencies"], use_proxy))
+
+    return upstream_tasks
diff --git a/test/test_util_taskcluster.py b/test/test_util_taskcluster.py
@@ -355,3 +355,150 @@ def test_list_task_group_incomplete_tasks(responses, root_url):
         },
     )
     assert list(tc.list_task_group_incomplete_tasks(tgid)) == ["1", "2", "3"]
+
+
+def test_get_ancestors(responses, root_url):
+    base_url = f"{root_url}/api/queue/v1/task"
+    responses.add(
+        responses.GET,
+        f"{base_url}/fff",
+        json={
+            "dependencies": ["eee", "ddd"],
+            "metadata": {
+                "name": "task-fff",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/eee",
+        json={
+            "dependencies": [],
+            "metadata": {
+                "name": "task-eee",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/ddd",
+        json={
+            "dependencies": ["ccc"],
+            "metadata": {
+                "name": "task-ddd",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/ccc",
+        json={
+            "dependencies": [],
+            "metadata": {
+                "name": "task-ccc",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/bbb",
+        json={
+            "dependencies": ["aaa"],
+            "metadata": {
+                "name": "task-bbb",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/aaa",
+        json={
+            "dependencies": [],
+            "metadata": {
+                "name": "task-aaa",
+            },
+        },
+    )
+
+    got = tc.get_ancestors(["bbb", "fff"])
+    expected = {
+        "task-aaa": "aaa",
+        "task-ccc": "ccc",
+        "task-ddd": "ddd",
+        "task-eee": "eee",
+    }
+    assert got == expected, f"got: {got}, expected: {expected}"
+
+
+def test_get_ancestors_string(responses, root_url):
+    base_url = f"{root_url}/api/queue/v1/task"
+    responses.add(
+        responses.GET,
+        f"{base_url}/fff",
+        json={
+            "dependencies": ["eee", "ddd"],
+            "metadata": {
+                "name": "task-fff",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/eee",
+        json={
+            "dependencies": [],
+            "metadata": {
+                "name": "task-eee",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/ddd",
+        json={
+            "dependencies": ["ccc", "bbb"],
+            "metadata": {
+                "name": "task-ddd",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/ccc",
+        json={
+            "dependencies": ["aaa"],
+            "metadata": {
+                "name": "task-ccc",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/bbb",
+        json={
+            "dependencies": [],
+            "metadata": {
+                "name": "task-bbb",
+            },
+        },
+    )
+    responses.add(
+        responses.GET,
+        f"{base_url}/aaa",
+        json={
+            "dependencies": [],
+            "metadata": {
+                "name": "task-aaa",
+            },
+        },
+    )
+
+    got = tc.get_ancestors("fff")
+    expected = {
+        "task-aaa": "aaa",
+        "task-bbb": "bbb",
+        "task-ccc": "ccc",
+        "task-ddd": "ddd",
+        "task-eee": "eee",
+    }
+    assert got == expected, f"got: {got}, expected: {expected}"