From 26bf989053d00539c4da463e40eadf2f1ea07d0b Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Fri, 5 Sep 2025 13:33:44 -0400
Subject: [PATCH 1/2] ci: add custom parameter for pull request

This will allow us to index tasks by pull request which can help testing
changes in external repos.
---
 .../self_taskgraph/custom_parameters.py       | 23 +++++++++++++++++++
 .../main-repo-pull-request-untrusted.yml      |  1 +
 .../test/params/main-repo-pull-request.yml    |  1 +
 3 files changed, 25 insertions(+)

diff --git a/taskcluster/self_taskgraph/custom_parameters.py b/taskcluster/self_taskgraph/custom_parameters.py
index 62d62c342..ae5b84c24 100644
--- a/taskcluster/self_taskgraph/custom_parameters.py
+++ b/taskcluster/self_taskgraph/custom_parameters.py
@@ -2,7 +2,30 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
+import os
+
+from voluptuous import All, Any, Range, Required
+
+from taskgraph.parameters import extend_parameters_schema
+
+
+def get_defaults(repo_root):
+    return {
+        "pull_request_number": None,
+    }
+
+
+extend_parameters_schema(
+    {
+        Required("pull_request_number"): Any(All(int, Range(min=1)), None),
+    },
+    defaults_fn=get_defaults,
+)
+
 
 def decision_parameters(graph_config, parameters):
     if parameters["tasks_for"] == "github-release":
         parameters["target_tasks_method"] = "release"
+
+    pr_number = os.environ.get("TASKGRAPH_PULL_REQUEST_NUMBER", None)
+    parameters["pull_request_number"] = None if pr_number is None else int(pr_number)
diff --git a/taskcluster/test/params/main-repo-pull-request-untrusted.yml b/taskcluster/test/params/main-repo-pull-request-untrusted.yml
index c48221c38..16ef26edb 100644
--- a/taskcluster/test/params/main-repo-pull-request-untrusted.yml
+++ b/taskcluster/test/params/main-repo-pull-request-untrusted.yml
@@ -17,6 +17,7 @@ moz_build_date: '20220421203159'
 optimize_target_tasks: true
 owner: user@example.com
 project: taskgraph
+pull_request_number: 123
 pushdate: 0
 pushlog_id: '0'
 repository_type: git
diff --git a/taskcluster/test/params/main-repo-pull-request.yml b/taskcluster/test/params/main-repo-pull-request.yml
index 1416c7d1d..a3b25d484 100755
--- a/taskcluster/test/params/main-repo-pull-request.yml
+++ b/taskcluster/test/params/main-repo-pull-request.yml
@@ -17,6 +17,7 @@ moz_build_date: '20220421203159'
 optimize_target_tasks: true
 owner: user@example.com
 project: taskgraph
+pull_request_number: 123
 pushdate: 0
 pushlog_id: '0'
 repository_type: git

From b2831c268a4947d7678eae57af523d1655058719 Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Fri, 5 Sep 2025 13:33:44 -0400
Subject: [PATCH 2/2] ci: add route to index docker-images by PR number

---
 taskcluster/kinds/docker-image/kind.yml       |  1 +
 .../self_taskgraph/transforms/add_pr_route.py | 31 +++++++++++++++++++
 2 files changed, 32 insertions(+)
 create mode 100644 taskcluster/self_taskgraph/transforms/add_pr_route.py

diff --git a/taskcluster/kinds/docker-image/kind.yml b/taskcluster/kinds/docker-image/kind.yml
index 06914ebf2..f53ab3853 100644
--- a/taskcluster/kinds/docker-image/kind.yml
+++ b/taskcluster/kinds/docker-image/kind.yml
@@ -10,6 +10,7 @@ loader: taskgraph.loader.transform:loader
 transforms:
     - taskgraph.transforms.docker_image:transforms
     - taskgraph.transforms.cached_tasks:transforms
+    - self_taskgraph.transforms.add_pr_route
     - taskgraph.transforms.task:transforms
 
 # make a task for each docker-image we might want.  For the moment, since we
diff --git a/taskcluster/self_taskgraph/transforms/add_pr_route.py b/taskcluster/self_taskgraph/transforms/add_pr_route.py
new file mode 100644
index 000000000..0f16d1497
--- /dev/null
+++ b/taskcluster/self_taskgraph/transforms/add_pr_route.py
@@ -0,0 +1,31 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from taskgraph.transforms.base import TransformSequence
+
+transforms = TransformSequence()
+
+
+@transforms.add
+def add_pr_route(config, tasks):
+    """Adds pull request index routes when applicable."""
+    if not (pr_number := config.params.get("pull_request_number")):
+        yield from tasks
+        return
+
+    PR_ROUTE = (
+        "index.{trust-domain}.v2.{project}-pr.{pr-number}.latest.{kind}.{task-name}"
+    )
+    subs = {
+        "trust-domain": config.graph_config["trust-domain"],
+        "project": config.params["project"],
+        "pr-number": pr_number,
+        "kind": config.kind,
+    }
+
+    for task in tasks:
+        subs["task-name"] = task.get("name") or task["label"][len(config.kind) + 1 :]
+        routes = task.setdefault("routes", [])
+        routes.append(PR_ROUTE.format(**subs))
+        yield task