From 5ec72c92917b583b8510653eafb2ba531f41a218 Mon Sep 17 00:00:00 2001 From: Ben Hearsum Date: Mon, 9 Dec 2024 20:54:12 -0500 Subject: [PATCH] feat: add support for fetching upstream tasks in find_tasks --- .../transforms/integration_test.py | 3 ++- .../fxci_config_taskgraph/util/integration.py | 26 +++++++++++++++++-- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/taskcluster/fxci_config_taskgraph/transforms/integration_test.py b/taskcluster/fxci_config_taskgraph/transforms/integration_test.py index 25fa9206..13e07b8d 100644 --- a/taskcluster/fxci_config_taskgraph/transforms/integration_test.py +++ b/taskcluster/fxci_config_taskgraph/transforms/integration_test.py @@ -197,5 +197,6 @@ def schedule_tasks_at_index(config, tasks): for task in tasks: for decision_index_path in task.pop("decision-index-paths"): - for task_def in find_tasks(decision_index_path): + include_deps = task.pop("include-deps", False) + for task_def in find_tasks(decision_index_path, include_deps): yield make_integration_test_description(task_def) diff --git a/taskcluster/fxci_config_taskgraph/util/integration.py b/taskcluster/fxci_config_taskgraph/util/integration.py index 0cd1661a..a4a8401b 100644 --- a/taskcluster/fxci_config_taskgraph/util/integration.py +++ b/taskcluster/fxci_config_taskgraph/util/integration.py @@ -3,15 +3,31 @@ # file, You can obtain one at http://mozilla.org/MPL/2.0/. from functools import cache -from typing import Any +import re +from typing import Any, Union +import os import requests import taskcluster +from taskgraph.util.taskcluster import get_ancestors as taskgraph_get_ancestors FIREFOXCI_ROOT_URL = "https://firefox-ci-tc.services.mozilla.com" STAGING_ROOT_URL = "https://stage.taskcluster.nonprod.cloudops.mozgcp.net" +def get_ancestors(task_ids: Union[list[str], str]) -> dict[str, str]: + # TODO: this is not ideal, but at the moment we don't have a better way + # to ensure that the upstream get_ancestors talks to the correct taskcluster + # instance. + orig = os.environ["TASKCLUSTER_ROOT_URL"] + try: + os.environ["TASKCLUSTER_ROOT_URL"] = FIREFOXCI_ROOT_URL + ret = taskgraph_get_ancestors(task_ids) + finally: + os.environ["TASKCLUSTER_ROOT_URL"] = orig + + return ret + @cache def get_taskcluster_client(service: str): options = {"rootUrl": FIREFOXCI_ROOT_URL} @@ -19,7 +35,7 @@ def get_taskcluster_client(service: str): @cache -def find_tasks(decision_index_path: str) -> list[dict[str, Any]]: +def find_tasks(decision_index_path: str, include_deps: list[re.Pattern] = []) -> list[dict[str, Any]]: """Find tasks targeted by the Decision task pointed to by `decision_index_path`.""" queue = get_taskcluster_client("queue") index = get_taskcluster_client("index") @@ -49,5 +65,11 @@ def find_tasks(decision_index_path: str) -> list[dict[str, Any]]: continue tasks.append(task["task"]) + # get_ancestors can be expensive; don't run it unless we might actually + # use the results. + if include_deps: + for label, task_id in get_ancestors(task["task_id"]).items(): + if any([pat.match(label) for pat in include_deps]): + tasks.append(queue.task(task_id)) return tasks