Add support for specific Docker tag in spawn_data_pipeline.py (#553)

* Revert "Revert "Add support for specific Docker tag in spawn_data_pipeline.py (#489)" (#499)"

This reverts commit 249ed40eb6.

* Ignore task with a tagged docker image

* Restrict Docker tag update to bugbug related images

infra/spawn_data_pipeline.py

@@ -69,6 +69,8 @@ def main():
 
     id_mapping = {}
 
+    docker_tag = os.getenv("TAG", None)
+
     # First pass, do the template rendering and dependencies resolution
     tasks = []
 
@@ -105,6 +107,30 @@ def main():
 
         payload["dependencies"] = new_dependencies
 
+        # Override the Docker image tag if needed
+        if docker_tag:
+            base_image = payload["payload"]["image"]
+
+            if base_image.startswith("mozilla/bugbug"):
+                splitted_image = base_image.rsplit(":", 1)
+
+                # If we have already a Docker tag
+                if len(splitted_image) > 1:
+                    docker_tag = splitted_image
+
+                    if docker_tag != "latest":
+                        msg = (
+                            "Don't update image {} for task {} as it already has a tag"
+                        )
+                        print(msg.format(base_image, task_internal_id))
+
+                tagless_image = splitted_image[0]
+
+                new_image = "{}:{}".format(tagless_image, docker_tag)
+                msg = "Updating image for task {} to {}"
+                print(msg.format(task_internal_id, new_image))
+                payload["payload"]["image"] = new_image
+
         tasks.append((task_id, payload))
 
     # Now sends them