feat: add support for custom context to task-context (#971)

* refactor: split paramatrized task_context tests into separate ones

This will make adding tests for `from-custom` context a bit more sane. It also corrects a small weirdness where `EXPECTED_DESCRIPTION` is a global, but it's actually just the expected value for one test.

* 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&regexp=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
```

Author: bhearsum

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.

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,18 @@ 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 []:
+            if name not in CUSTOM_CONTEXT_MAP:
+                raise ValueError(f"no provider found for custom context '{name}'")
+
+            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", {}))

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

test/data/task_context.yml

@@ -1,4 +1,5 @@
 file: file
+file_and_custom: file_and_custom
 direct_and_file: shouldn't be used
 file_and_param: shouldn't be used
 direct_file_and_param: shouldn't be used