ci: index docker-images by pull request number

src/taskgraph/transforms/docker_image.py

@@ -256,3 +256,31 @@ def fill_template(config, tasks):
             }
 
         yield taskdesc
+
+
+@transforms.add
+def add_pr_route(config, tasks):
+    """Taskgraph builds several docker-images that get used externally.
+
+    Indexing these images by pull request number, allows us to easily test them
+    in external repos using the `indexed-image` image definition type.
+    """
+    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.{product}.{task-name}"
+    )
+    subs = {
+        "trust-domain": config.graph_config["trust-domain"],
+        "project": config.params["project"],
+        "pr-number": pr_number,
+        "product": "docker-image",
+    }
+
+    for task in tasks:
+        subs["task-name"] = task["attributes"]["image_name"]
+        routes = task.setdefault("routes", [])
+        routes.append(PR_ROUTE.format(**subs))
+        yield task

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)

taskcluster/test/params/main-repo-pull-request-untrusted.yml

@@ -17,6 +17,7 @@ moz_build_date: '20220421203159'
 optimize_target_tasks: true
 owner: [email protected]
 project: taskgraph
+pull_request_number: 123
 pushdate: 0
 pushlog_id: '0'
 repository_type: git

taskcluster/test/params/main-repo-pull-request.yml

@@ -17,6 +17,7 @@ moz_build_date: '20220421203159'
 optimize_target_tasks: true
 owner: [email protected]
 project: taskgraph
+pull_request_number: 123
 pushdate: 0
 pushlog_id: '0'
 repository_type: git

test/test_transform_docker_image.py

@@ -43,6 +43,22 @@
             },
             id="symbol",
         ),
+        pytest.param(
+            {},
+            {"pull_request_number": "123"},
+            {
+                "routes": [
+                    "index.test-domain.v2.some-project-pr.123.latest.docker-image.fake-name"
+                ]
+            },
+            id="pr_route",
+        ),
+        pytest.param(
+            {},
+            {},
+            {},
+            id="no_pr_route_without_pr_number",
+        ),
     ),
 )
 @unittest.mock.patch(