From d580540d1b74fbe725f5add808a7ae82cc34ec77 Mon Sep 17 00:00:00 2001 From: Ben Hearsum Date: Wed, 3 Jun 2026 06:52:53 -0400 Subject: [PATCH] feat: add support for custom context to task-context This is somewhat of a continuation to #968 in that it would allow for more `resolve_keyed_by` calls to be replaced by `task-context`. As a concrete example, we use [`by-release-level` in many places in gecko](https://searchfox.org/firefox-main/search?q=by-release-level&path=&case=false®exp=false). Its value is [derived from parameters and some constants](https://searchfox.org/firefox-main/source/taskcluster/gecko_taskgraph/util/attributes.py#157). With this patch, we could register it as custom context and replace a whole bunch of `resolve_keyed_by` calls in transforms, eg: In some python code that's called at taskgraph registration: ``` @custom_context("release-level") def release_level_context(config, task): """Provide ``release-level`` ("production" or "staging"), derived from parameters, so ``by-release-level`` fields can be resolved declaratively.""" return {"release-level": release_level(config.params)} ``` And in various kinds we'd end up with things like: ``` transforms: - taskgraph.transforms.task_context task-defaults: task-context: from-custom: - release-level substitution-fields: - scopes scopes: by-release-level: production: - project:releng:ship-it:server:production - project:releng:ship-it:action:create-new-release staging: - project:releng:ship-it:server:staging - project:releng:ship-it:action:create-new-release ``` --- docs/reference/transforms/task_context.rst | 111 +++++++++++++++------ src/taskgraph/transforms/task_context.py | 13 +++ src/taskgraph/util/task_context.py | 20 ++++ test/test_transform_task_context.py | 33 ++++++ 4 files changed, 147 insertions(+), 30 deletions(-) create mode 100644 src/taskgraph/util/task_context.py diff --git a/docs/reference/transforms/task_context.rst b/docs/reference/transforms/task_context.rst index 9244c6098..f11c2354c 100644 --- a/docs/reference/transforms/task_context.rst +++ b/docs/reference/transforms/task_context.rst @@ -98,35 +98,6 @@ with the this ``kind``: ...the description will bring in the value ``foo`` from the parameters if present, or ``default`` otherwise. -from-file -~~~~~~~~~ - -Context may also be provided from a defined yaml file. The provided file -should usually only contain top level keys and values (eg: nested objects -will not be interpolated - they will be substituted as text representations -of the object). - -For example, with this kind definition: - -.. code-block:: yaml - - tasks: - build: - description: my description {foo} - task-context: - from-file: some_file.yaml - substitution-fields: - - description - -And this in ``some_file.yaml``: - -.. code-block:: yaml - - foo: from a file - -...description will end up with "my description from a file". - - from-object ~~~~~~~~~~~ @@ -158,6 +129,86 @@ For example: This will give build1 and build2 descriptions with their ``extra_desc`` included while allowing them to share the rest of their task definition. +from-custom +~~~~~~~~~~~ + +Context may be provided by custom providers, which must be registered prior +to this transform being run. This allows the creation of context that is +derived from parameters, code constants, or anything else accessible to +taskgraph. + +For example, you may have a custom context handler set-up in +``taskcluster/my_taskgraph/custom_context.py``: + +.. code-block:: python + + NON_PRODUCTION_BRANCHES = ["maple", "pine"] + + @custom_context("release-level") + def release_level_context(config, task): + # Despite being level 3, some branches are not truly considered "production" + # in the sense of creating releases that ship to users. + if config.params["level"] == "1" or config.params["project"] in NON_PRODUCTION_BRANCHES: + return "staging" + return "production" + + +In your ``register`` function you will need to ensure you import this module. +Eg: ``taskcluster/my_taskgraph/__init__.py``: + +.. code-block:: python + + def register(graph_config): + from my_taskgraph import custom_contexts # trigger custom task-context registration + + +Now you can use ``release-level`` as context in a kind: + +.. code-block:: yaml + + task-defaults: + task-context: + from-custom: + - release-level + substitution-fields: + - scopes + + scopes: + by-release-level: + staging: + - secrets:get:staging_creds + production: + - secrets:get:production_creds + + +from-file +~~~~~~~~~ + +Context may also be provided from a defined yaml file. The provided file +should usually only contain top level keys and values (eg: nested objects +will not be interpolated - they will be substituted as text representations +of the object). + +For example, with this kind definition: + +.. code-block:: yaml + + tasks: + build: + description: my description {foo} + task-context: + from-file: some_file.yaml + substitution-fields: + - description + +And this in ``some_file.yaml``: + +.. code-block:: yaml + + foo: from a file + +...description will end up with "my description from a file". + Implicit Context ~~~~~~~~~~~~~~~~ @@ -185,7 +236,7 @@ Keys will be resolved on ``substitution-fields`` first, then substitution will be performed on the resolved value. If the same key is found in multiple places the order of precedence is as -follows: ``from-parameters``, ``from-object`` keys, ``from-file`` and finally +follows: ``from-parameters``, ``from-object`` keys, ``from-custom`` providers, ``from-file`` and finally implicit context. That is to say: parameters will always override anything else. diff --git a/src/taskgraph/transforms/task_context.py b/src/taskgraph/transforms/task_context.py index 4304d42c4..3d01d87aa 100644 --- a/src/taskgraph/transforms/task_context.py +++ b/src/taskgraph/transforms/task_context.py @@ -2,6 +2,7 @@ from taskgraph.transforms.base import TransformSequence from taskgraph.util.schema import Schema, resolve_keyed_by +from taskgraph.util.task_context import CUSTOM_CONTEXT_MAP from taskgraph.util.templates import deep_get, substitute_task_fields from taskgraph.util.yaml import load_yaml @@ -28,6 +29,12 @@ class TaskContextConfig(Schema): from_file: Optional[str] = None # Key/value pairs to be used as task context from_object: Optional[object] = None + # Retrieve task context values from registered custom providers. Each + # entry is the name of a provider registered via + # ``taskgraph.util.task_context.custom_context``. Providers are called + # with the ``TransformConfig`` and the task being rendered, and must + # return a dict of key/value pairs to add to the context. + from_custom: Optional[list[str]] = None resolve_keyed_by_options: Optional[ResolveKeyedByConfigOptions] = None @@ -51,6 +58,7 @@ class TaskContextSchema(Schema, forbid_unknown_fields=False, kw_only=True): # is as follows: # - Parameters # - `from-object` keys + # - Custom providers (`from-custom`) # - File # # That is to say: parameters will always override anything else. @@ -89,10 +97,15 @@ def render_task(config, tasks): if from_file: file_context = load_yaml(from_file) + custom_context = {} + for name in sub_config.pop("from-custom", None) or []: + custom_context.update(CUSTOM_CONTEXT_MAP[name](config, task)) + fields = sub_config.pop("substitution-fields") subs = {} subs.update(file_context) + subs.update(custom_context) # We've popped away the configuration; everything left in `sub_config` is # substitution key/value pairs. subs.update(sub_config.pop("from-object", {})) diff --git a/src/taskgraph/util/task_context.py b/src/taskgraph/util/task_context.py new file mode 100644 index 000000000..2b9c8b3e5 --- /dev/null +++ b/src/taskgraph/util/task_context.py @@ -0,0 +1,20 @@ +# 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/. + +# Define a collection of custom task-context providers. +# Note: this is stored here instead of where it is used in the +# `task_context` transform to give consumers a chance to register their own +# providers before the `task_context` schema is created. +CUSTOM_CONTEXT_MAP = {} + + +def custom_context(name): + def wrapper(func): + assert name not in CUSTOM_CONTEXT_MAP, ( + f"duplicate custom_context function name {name} ({func} and {CUSTOM_CONTEXT_MAP[name]})" + ) + CUSTOM_CONTEXT_MAP[name] = func + return func + + return wrapper diff --git a/test/test_transform_task_context.py b/test/test_transform_task_context.py index 8f02ebaed..f125e95f3 100644 --- a/test/test_transform_task_context.py +++ b/test/test_transform_task_context.py @@ -11,6 +11,7 @@ from taskgraph.transforms import task_context from taskgraph.transforms.base import TransformConfig +from taskgraph.util.task_context import CUSTOM_CONTEXT_MAP, custom_context here = os.path.abspath(os.path.dirname(__file__)) @@ -220,3 +221,35 @@ def test_resolve_keyed_by_missing_context_with_defer(run_transform, graph_config "default": "also-should-not-resolve", }, } + + +def test_resolve_keyed_by_custom_context(run_transform, graph_config): + try: + + @custom_context("colour") + def _colour(config, task): + return {"colour": "blue"} + + task = { + "name": "fake-task-name", + "description": { + "by-colour": { + "blue": "i am a blue task!", + "default": " i am a default task", + }, + }, + "task-context": { + "from-custom": ["colour"], + "substitution-fields": ["description"], + }, + } + config = make_config(graph_config) + + task = run_transform(task_context.transforms, task, config=config)[0] + pprint(task, indent=2) + + assert task["description"] == "i am a blue task!" + finally: + # ensure `CUSTOM_CONTEXT_MAP` is reset at the end of the test + # regardless of result + CUSTOM_CONTEXT_MAP.pop("colour", None)