diff --git a/src/taskgraph/util/taskcluster.py b/src/taskgraph/util/taskcluster.py index 9946f2912..e08432cf6 100644 --- 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 index 84c233089..7b29515a1 100644 --- 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}"