feat: add config-driven release_level utility

src/mozilla_taskgraph/__init__.py

@@ -7,6 +7,8 @@
 from taskgraph.config import validate_graph_config
 from taskgraph.util import schema
 
+from mozilla_taskgraph.parameters import register_parameters
+
 # Schemas for YAML files should use dashed identifiers by default. If there are
 # components of the schema for which there is a good reason to use another format,
 # exceptions can be added here.
@@ -26,6 +28,8 @@ def register(graph_config):
             "worker_types",
         ]
     )
+
+    register_parameters()
     validate_graph_config(graph_config._config)
 
 

src/mozilla_taskgraph/config.py

@@ -1,5 +1,5 @@
 from textwrap import dedent
-from typing import Optional
+from typing import Optional, Union
 
 from taskgraph import config as tg
 from taskgraph.util.schema import Schema
@@ -26,6 +26,11 @@ class MozillaGraphConfigSchema(tg.graph_config_schema):
         # string to use for release tasks.
         # Defaults to ``mozilla_taskgraph.version:default_parser``.
         version_parser: Optional[str] = None
+        # Mapping of project to the branches that should be considered
+        # "production" releases. A value of ``True`` means all branches of the
+        # project are release branches, while a list restricts releases to the
+        # named branches. Consumed by ``mozilla_taskgraph.parameters``.
+        release_branches: Optional[dict[str, Union[bool, list[str]]]] = None
 
 else:
     # Legacy voluptuous-based graph_config_schema (e.g. gecko_taskgraph override).
@@ -48,6 +53,15 @@ class MozillaGraphConfigSchema(tg.graph_config_schema):
                     Defaults to ``mozilla_taskgraph.version:default_parser``.
                     """.lstrip()),
             ): str,
+            Vol_Optional(
+                "release-branches",
+                description=dedent("""
+                    Mapping of project to the branches that should be considered
+                    "production" releases. A value of ``True`` means all branches
+                    of the project are release branches, while a list restricts
+                    releases to the named branches.
+                    """.lstrip()),
+            ): {str: object},
         }
     )
 

src/mozilla_taskgraph/parameters.py

@@ -0,0 +1,38 @@
+# 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 typing import Optional, Union
+
+from taskgraph.parameters import extend_parameters_schema
+from taskgraph.util.schema import Schema
+
+
+class MozillaParametersSchema(Schema, kw_only=True):
+    # Branches of the current project that are considered "production"
+    # releases. Resolved from the graph config's `release-branches` mapping by
+    # `set_release_branches`. `True` means all branches of the project are
+    # release branches, a list restricts releases to the named branches, and
+    # `None` means the project has no release branches.
+    release_branches: Optional[Union[bool, list[str]]] = None
+
+
+def get_defaults(repo_root=None):
+    return {
+        "release_branches": None,
+    }
+
+
+def register_parameters():
+    extend_parameters_schema(MozillaParametersSchema, defaults_fn=get_defaults)
+
+
+def set_release_branches(graph_config, parameters):
+    """Resolve the current project's release branches from the graph config.
+
+    Projects should call this from their `decision-parameters` function so that
+    `release_branches` is persisted into `parameters.yml` and available to
+    consumers that only have a `Parameters` object (e.g. action tasks).
+    """
+    mapping = graph_config._config.get("release-branches") or {}
+    parameters["release_branches"] = mapping.get(parameters["project"])

src/mozilla_taskgraph/util/attributes.py

@@ -0,0 +1,34 @@
+# 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/.
+
+import re
+
+
+def release_level(params):
+    """Whether this is a production release or not.
+
+    The set of branches considered "production" is project specific and comes
+    from the ``release_branches`` parameter, which is resolved at decision time
+    from the graph config's ``release-branches`` mapping. A value of ``True``
+    means every branch of the project is a release branch (the model used by
+    Mercurial based projects), while a list restricts releases to the named
+    branches.
+
+    :return str: One of "production" or "staging".
+    """
+    if params["level"] != "3":
+        return "staging"
+
+    branches = params.get("release_branches")
+    if not branches:
+        return "staging"
+
+    if branches is True:
+        return "production"
+
+    m = re.match(r"refs/heads/(\S+)$", params["head_ref"])
+    if m is not None and m.group(1) in branches:
+        return "production"
+
+    return "staging"

test/test_parameters.py

@@ -0,0 +1,34 @@
+# 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/.
+
+import pytest
+
+from mozilla_taskgraph.parameters import set_release_branches
+
+
+@pytest.mark.parametrize(
+    "release_branches,project,expected",
+    (
+        # No `release-branches` config at all.
+        (None, "some-project", None),
+        # Project not listed.
+        ({"firefox": ["main"]}, "autoland", None),
+        # Whole project is a release project.
+        ({"mozilla-central": True}, "mozilla-central", True),
+        # Project mapped to a list of branches.
+        ({"firefox": ["main", "beta"]}, "firefox", ["main", "beta"]),
+    ),
+)
+def test_set_release_branches(
+    make_graph_config, parameters, release_branches, project, expected
+):
+    extra_config = {}
+    if release_branches is not None:
+        extra_config["release-branches"] = release_branches
+
+    graph_config = make_graph_config(extra_config=extra_config)
+    parameters["project"] = project
+
+    set_release_branches(graph_config, parameters)
+    assert parameters["release_branches"] == expected

test/util/test_attributes.py

@@ -0,0 +1,58 @@
+# 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/.
+
+import pytest
+
+from mozilla_taskgraph.util.attributes import release_level
+
+FIREFOX_BRANCHES = ["main", "beta", "release", "esr140"]
+
+
+@pytest.mark.parametrize(
+    "params,expected",
+    (
+        # Not level 3 -> always staging, regardless of branch.
+        ({"level": "1", "release_branches": True}, "staging"),
+        (
+            {
+                "level": "1",
+                "release_branches": FIREFOX_BRANCHES,
+                "head_ref": "refs/heads/beta",
+            },
+            "staging",
+        ),
+        # No release branches for the project -> staging (e.g. autoland).
+        ({"level": "3", "release_branches": None}, "staging"),
+        # Whole project is a release project (Mercurial model).
+        ({"level": "3", "release_branches": True}, "production"),
+        # Git monorepo model: only listed branches are production.
+        (
+            {
+                "level": "3",
+                "release_branches": FIREFOX_BRANCHES,
+                "head_ref": "refs/heads/beta",
+            },
+            "production",
+        ),
+        (
+            {
+                "level": "3",
+                "release_branches": FIREFOX_BRANCHES,
+                "head_ref": "refs/heads/test",
+            },
+            "staging",
+        ),
+        # Only refs/heads/* match, not tags.
+        (
+            {
+                "level": "3",
+                "release_branches": FIREFOX_BRANCHES,
+                "head_ref": "refs/tags/beta",
+            },
+            "staging",
+        ),
+    ),
+)
+def test_release_level(params, expected):
+    assert release_level(params) == expected