From d64ad0967bd3245b370e7c75c2db6771d702f61c Mon Sep 17 00:00:00 2001
From: Tom Prince <mozilla@hocat.ca>
Date: Tue, 7 Jul 2020 19:56:11 +0000
Subject: [PATCH] Bug 1626058: Add support for writting out artifacts from
 transforms in the decision task; r=Callek,ahal

To support using kaniko[1] for building images, we need to generate the docker
contexts in a seprate task from the docker-image task. Since we use the hash of
the context as the cached-task digest, we generate the context in the decision
task already, so this adds a way to write that out to be used by downstream
tasks.

[1] https://github.com/GoogleContainerTools/kaniko

Differential Revision: https://phabricator.services.mozilla.com/D77839
---
 taskcluster/taskgraph/decision.py               |  1 +
 taskcluster/taskgraph/generator.py              | 17 +++++++++++++----
 taskcluster/taskgraph/test/test_generator.py    |  4 ++--
 .../taskgraph/test/test_transforms_job.py       |  2 +-
 taskcluster/taskgraph/transforms/base.py        |  3 +++
 5 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/taskcluster/taskgraph/decision.py b/taskcluster/taskgraph/decision.py
index bc53417c8091..912cbf40e56b 100644
--- a/taskcluster/taskgraph/decision.py
+++ b/taskcluster/taskgraph/decision.py
@@ -229,6 +229,7 @@ def taskgraph_decision(options, parameters=None):
         root_dir=options.get('root'),
         parameters=parameters,
         decision_task_id=decision_task_id,
+        write_artifacts=True,
     )
 
     # write out the parameters used to generate this graph
diff --git a/taskcluster/taskgraph/generator.py b/taskcluster/taskgraph/generator.py
index 00586e4bcccf..c02885e8d994 100644
--- a/taskcluster/taskgraph/generator.py
+++ b/taskcluster/taskgraph/generator.py
@@ -49,7 +49,7 @@ class Kind(object):
             raise KeyError("{!r} does not define `loader`".format(self.path))
         return find_object(loader)
 
-    def load_tasks(self, parameters, loaded_tasks):
+    def load_tasks(self, parameters, loaded_tasks, write_artifacts):
         loader = self._get_loader()
         config = copy.deepcopy(self.config)
 
@@ -66,7 +66,8 @@ class Kind(object):
 
         # perform the transformations on the loaded inputs
         trans_config = TransformConfig(self.name, self.path, config, parameters,
-                                       kind_dependencies_tasks, self.graph_config)
+                                       kind_dependencies_tasks, self.graph_config,
+                                       write_artifacts=write_artifacts)
         tasks = [Task(self.name,
                       label=task_dict['label'],
                       attributes=task_dict['attributes'],
@@ -108,7 +109,12 @@ class TaskGraphGenerator(object):
     # circuit generation of the entire graph by never completing the generator.
 
     def __init__(
-            self, root_dir, parameters, decision_task_id="<decision-task>", target_kind=None,
+        self,
+        root_dir,
+        parameters,
+        decision_task_id="DECISION-TASK",
+        write_artifacts=False,
+        target_kind=None,
     ):
         """
         @param root_dir: root directory, with subdirectories for each kind
@@ -122,6 +128,7 @@ class TaskGraphGenerator(object):
         self._parameters = parameters
         self._target_kind = target_kind
         self._decision_task_id = decision_task_id
+        self._write_artifacts = write_artifacts
 
         # start the generator
         self._run = self._run()
@@ -267,7 +274,9 @@ class TaskGraphGenerator(object):
             logger.debug("Loading tasks for kind {}".format(kind_name))
             kind = kinds[kind_name]
             try:
-                new_tasks = kind.load_tasks(parameters, list(all_tasks.values()))
+                new_tasks = kind.load_tasks(
+                    parameters, list(all_tasks.values()), self._write_artifacts,
+                )
             except Exception:
                 logger.exception("Error loading tasks for kind {}:".format(kind_name))
                 raise
diff --git a/taskcluster/taskgraph/test/test_generator.py b/taskcluster/taskgraph/test/test_generator.py
index 23388bda1033..6a0a873f88c2 100644
--- a/taskcluster/taskgraph/test/test_generator.py
+++ b/taskcluster/taskgraph/test/test_generator.py
@@ -45,9 +45,9 @@ class FakeKind(Kind):
     def _get_loader(self):
         return fake_loader
 
-    def load_tasks(self, parameters, loaded_tasks):
+    def load_tasks(self, parameters, loaded_tasks, write_artifacts):
         FakeKind.loaded_kinds.append(self.name)
-        return super(FakeKind, self).load_tasks(parameters, loaded_tasks)
+        return super(FakeKind, self).load_tasks(parameters, loaded_tasks, write_artifacts)
 
 
 class WithFakeKind(TaskGraphGenerator):
diff --git a/taskcluster/taskgraph/test/test_transforms_job.py b/taskcluster/taskgraph/test/test_transforms_job.py
index d65a7833e646..575d719209c7 100644
--- a/taskcluster/taskgraph/test/test_transforms_job.py
+++ b/taskcluster/taskgraph/test/test_transforms_job.py
@@ -37,7 +37,7 @@ TASK_DEFAULTS = {
 @pytest.fixture(scope='module')
 def config():
     graph_config = load_graph_config(os.path.join(GECKO, 'taskcluster', 'ci'))
-    return TransformConfig('job_test', here, {}, {}, [], graph_config)
+    return TransformConfig('job_test', here, {}, {}, [], graph_config, write_artifacts=False)
 
 
 @pytest.fixture()
diff --git a/taskcluster/taskgraph/transforms/base.py b/taskcluster/taskgraph/transforms/base.py
index 53605f61b9c0..63c9fe75da58 100644
--- a/taskcluster/taskgraph/transforms/base.py
+++ b/taskcluster/taskgraph/transforms/base.py
@@ -38,6 +38,9 @@ class TransformConfig(object):
     # Global configuration of the taskgraph
     graph_config = attr.ib(type=GraphConfig)
 
+    # whether to write out artifacts for the decision task
+    write_artifacts = attr.ib(type=bool)
+
 
 @attr.s()
 class TransformSequence(object):